# Copyright (C) 2012 Victor Semionov
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#  * Redistributions of source code must retain the above copyright notice, this
#    list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import os

import math

import copy

import numpy as np

import model

import params
import core
import plot
import output
import unitconv

def _set_geom(rm, rb):
    old_geom = params.medium_radius, params.beam_radius
    try:
        params.medium_radius, params.beam_radius = rm, rb
        return old_geom
    except:
        params.medium_radius, params.beam_radius = old_geom
        raise

def _ref_signal_fluence(active_medium, (rho, phi), (integrator, amp), count_t, ref_pulse, lower_decay):
    pulse_count = params.train_pulse_count

    upper = np.vectorize(active_medium.initial_inversion.inversion)(rho, phi, amp.Z)
    lower = np.zeros(len(amp.Z))
    population = (upper, lower)

    amp._init_time(ref_pulse, count_t)
    input_density = np.vectorize(ref_pulse.density)(amp.T)

    pulse_fluences = np.empty(pulse_count)

    for pnum in range(pulse_count):
        density_out, population_final = amp.amplify(rho, phi, None, None, T=amp.T, initial_population=population, input_density=input_density)

        upper = np.copy(population_final[0])
        lower = population_final[1] * lower_decay
        population = (upper, lower)

        fluence_out = integrator.integrate(amp.T, density_out) * active_medium.light_speed
        pulse_fluences[pnum] = fluence_out

    fluence_out = pulse_fluences[::-1].sum()
    fluence_out = model.energy.energy(params.lasing_wavelen, fluence_out)

    return fluence_out

def compare_depop_models(dirname):
    filename = lambda name: os.path.join(dirname, name)

    if not params.ext_depop_models:
        return

    print output.div_line
    print "comparing depopulation models"

    active_medium = core.create_medium(None)
    pump_system = model.pump.PumpSystem(params.pump_wavelen, params.pump_duration, params.pump_power, params.pump_efficiency)
    data = []
    for depop_model_class in params.ext_depop_models:
        depop_model_label = depop_model_class.descr
        print depop_model_label
        depop_model = core.create_depop_model(active_medium, depop_model_class)
        inv = params.inverter_class(active_medium, pump_system, depop_model)
        inv.invert(params.inversion_rtol, params.inversion_min_count_t)
        depop_rate = np.vectorize(depop_model.rate)(inv.inversion) / active_medium.volume
        data.append((inv.T, inv.inversion, depop_rate, depop_model_class.descr, depop_model_class))
        ref_inversion = inv.inversion[-1]

        unitconv.print_result("population inversion [{}]: {}", ("cm^-3",), (ref_inversion,))
        unitconv.print_result("depopulation rate [{}]: {}", ("cm^-3 s^-1",), (depop_rate[-1],))

    if params.graphs:
        print output.status_writing
        dirname = os.path.join(dirname, output.models_rel_path)
        dirname = output.init_dir(dirname)
        data.sort(key = lambda x: x[1][-1], reverse=True)
        Ts, inversions, depop_rates, labels, depop_model_classes = zip(*data)
        plot.plot_data(filename("inversions_evo"), "Population Inversion Evolution", (Ts, None, None, output.t_pump_label), (inversions, None, None, output.inversion_abs_label), labels)
        pump_rate = pump_system.effective_pump_rate / active_medium.volume
        abs_rate_ylim = None #(0.0, pump_rate * 1.25)
        non_zero_Ts = [T[1:] for T in Ts]
        non_zero_inversions = [inversion[1:] for inversion in inversions]
        non_zero_rates = [depop_rate[1:] for depop_rate in depop_rates]
        rel_depop_rates = [depop_rate / inversion for depop_rate, inversion in zip(non_zero_rates, non_zero_inversions)]
        plot.plot_data(filename("depop_rates"), "Depopulation Rate", (inversions, None, None, output.inversion_abs_label), (depop_rates, None, abs_rate_ylim, output.rate_label), labels, yvals=[(pump_rate, "pump rate")])
        plot.plot_data(filename("depop_rates_alt"), "Depopulation Rate to Inversion Ratio", (non_zero_inversions, None, None, output.inversion_abs_label), (rel_depop_rates, None, None, output.rate_rel_label), labels)
        plot.plot_data(filename("depop_rates_evo"), "Depopulation Rate Evolution", (Ts, None, None, output.t_pump_label), (depop_rates, None, abs_rate_ylim, output.rate_label), labels, yvals=[(pump_rate, "pump rate")])
        plot.plot_data(filename("depop_rates_alt_evo"), "Depopulation Rate to Inversion Ratio Evolution", (non_zero_Ts, None, None, output.t_pump_label), (rel_depop_rates, None, None, output.rate_rel_label), labels)

        if params.ext_alt_depop_model not in depop_model_classes:
            return
        alt_model_idx = depop_model_classes.index(params.ext_alt_depop_model)
        alt_T = Ts[alt_model_idx]
        alt_inversion = inversions[alt_model_idx]
        altinvs = []
        aTs = []
        altinv_inversion_rdiffs = []
        for cls, T, inversion in zip(depop_model_classes, Ts, inversions):
            if cls is params.ext_alt_depop_model:
                continue
            uT = set(list(T) + list(alt_T))
            aT = np.array(sorted(list(uT)))
            altinv = np.interp(aT, alt_T, alt_inversion)[1:]
            inv = np.interp(aT, T, inversion)[1:]
            rdiff = np.fabs(inv - altinv) / np.fmin(inv, altinv)
            aTs.append(aT[1:])
            altinvs.append(altinv)
            altinv_inversion_rdiffs.append(rdiff)
        non_alt_labels = [label for i, label in enumerate(labels) if i != alt_model_idx]
        plot.plot_data(filename("inversions_rdiff_inv"), "Inversion Relative Difference", (altinvs, None, None, output.inversion_abs_label), (altinv_inversion_rdiffs, None, None, output.inversion_rdiff_label), non_alt_labels)
        plot.plot_data(filename("inversions_rdiff_evo"), "Inversion Relative Difference Evolution", (aTs, None, None, output.t_pump_label), (altinv_inversion_rdiffs, None, None, output.inversion_rdiff_label), non_alt_labels)

def _inversion_pump_dependence_task((i, j), (tau, pwr), active_medium, (depop_model1, depop_model2)):
    output.show_status((i, j), params.extended_status_strides, False)

    pump_system = model.pump.PumpSystem(params.pump_wavelen, tau, pwr, params.pump_efficiency)

    inv1 = params.inverter_class(active_medium, pump_system, depop_model1)
    inv2 = params.inverter_class(active_medium, pump_system, depop_model2)

    inversion1 = inv1.invert(params.inversion_rtol, params.inversion_min_count_t)
    inversion2 = inv2.invert(params.inversion_rtol, params.inversion_min_count_t)

    gain_coef1 = inversion1 * active_medium.doping_agent.xsection
    gain_coef2 = inversion2 * active_medium.doping_agent.xsection

    gain1 = math.exp(gain_coef1 * active_medium.length)
    gain2 = math.exp(gain_coef2 * active_medium.length)

    inversion_rdiff = abs((inversion1 - inversion2) / min(inversion1, inversion2))
    gain_rdiff = abs((gain1 - gain2) / min(gain1, gain2))

    return inversion1, inversion2, gain1, gain2, inversion_rdiff, gain_rdiff

def compute_inversion_pump_dependence(task_pool, dirname):
    filename = lambda name: os.path.join(dirname, name)

    print output.div_line
    print "computing inversion dependence on pumping parameters"

    active_medium = core.create_medium(None)

    count_tau = params.ext_opt_pump_resolution[0]
    count_pwr = params.ext_opt_pump_resolution[1]

    Tau = np.linspace(params.ext_opt_pump_duration[0], params.ext_opt_pump_duration[1], count_tau)
    Pwr = np.linspace(params.ext_opt_pump_power[0], params.ext_opt_pump_power[1], count_pwr)

    depop_model_class1 = params.depop_model_class
    depop_model_class2 = params.ext_alt_depop_model

    depop_model1 = core.create_depop_model(active_medium, params.depop_model_class)
    depop_model2 = core.create_depop_model(active_medium, params.ext_alt_depop_model)

    inversions1, inversions2, gains1, gains2, inversion_rdiffs, gain_rdiffs = task_pool.parallel_task(_inversion_pump_dependence_task, (Tau, Pwr), (), (active_medium, (depop_model1, depop_model2)))

    output.show_status((count_tau, count_pwr), params.extended_status_strides, True)

    pump_energies = np.prod(np.array(np.meshgrid(Tau, Pwr)), axis=0).T
    stored_energies1 = model.energy.energy(params.lasing_wavelen, inversions1 * active_medium.volume)
    stored_energies2 = model.energy.energy(params.lasing_wavelen, inversions2 * active_medium.volume)

    if params.graphs:
        print output.status_writing
        dirname = os.path.join(dirname, output.opt_pump_rel_path)
        inversion_rdiff_max = params.ext_opt_inversion_rdiff_max
        zlim = None #(0.0, inversion_rdiff_max)
        depop_contours = [inversion_rdiff_max]
        ref_pump_energy = params.pump_duration * params.pump_power
        ref_pump_contours = [(pump_energies.T, ref_pump_energy, "const. pump energy")]
        graph_types = [
            (dirname, Pwr, output.pump_power_label),
            (os.path.join(dirname, output.alt_plot_rel_path), Pwr * params.pump_efficiency / active_medium.volume, output.eff_power_density_label),
        ]
        for dirname, Y, ylabel in graph_types:
            dirname = output.init_dir(dirname)
            plot.plot_color(filename("energy_pump"), "Pump Energy", (Tau, None, None, output. pump_duration_label), (Y, None, None, ylabel), (pump_energies.T, None, output.energy_abs_pump_label), params.out_num_auto_contours)
            plot.plot_color(filename("inversion"), "Inversion (%s)" % depop_model_class1.descr, (Tau, None, None, output.pump_duration_label), (Y, None, None, ylabel), (inversions1.T, None, output.inversion_abs_label), params.out_num_auto_contours, extra_contours=ref_pump_contours)
            plot.plot_color(filename("inversion_alt"), "Inversion (%s)" % depop_model_class2.descr, (Tau, None, None, output.pump_duration_label), (Y, None, None, ylabel), (inversions2.T, None, output.inversion_abs_label), params.out_num_auto_contours, extra_contours=ref_pump_contours)
            plot.plot_color(filename("ss_gain"), "Small Signal Gain (%s)" % depop_model_class1.descr, (Tau, None, None, output.pump_duration_label), (Y, None, None, ylabel), (gains1.T, None, output.gain_label), params.out_num_auto_contours, extra_contours=ref_pump_contours)
            plot.plot_color(filename("ss_gain_alt"), "Small Signal Gain (%s)" % depop_model_class2.descr, (Tau, None, None, output.pump_duration_label), (Y, None, None, ylabel), (gains2.T, None, output.gain_label), params.out_num_auto_contours, extra_contours=ref_pump_contours)
            plot.plot_color(filename("energy_stored"), "Stored Energy (%s)" % depop_model_class1.descr, (Tau, None, None, output.pump_duration_label), (Y, None, None, ylabel), (stored_energies1.T, None, output.energy_abs_stored_label), params.out_num_auto_contours, extra_contours=ref_pump_contours)
            plot.plot_color(filename("energy_stored_alt"), "Stored Energy (%s)" % depop_model_class2.descr, (Tau, None, None, output.pump_duration_label), (Y, None, None, ylabel), (stored_energies2.T, None, output.energy_abs_stored_label), params.out_num_auto_contours, extra_contours=ref_pump_contours)
            plot.plot_color(filename("inversion_rdiff"), "Inversion Relative Difference", (Tau, None, None, output.pump_duration_label), (Y, None, None, ylabel), (inversion_rdiffs.T, zlim, output.inversion_rdiff_label), params.out_num_auto_contours, depop_contours)
            plot.plot_color(filename("ss_gain_rdiff"), "Small Signal Gain Relative Difference", (Tau, None, None, output.pump_duration_label), (Y, None, None, ylabel), (gain_rdiffs.T, zlim, output.gain_rdiff_label), params.out_num_auto_contours, depop_contours)

    return inversions1, inversion_rdiffs

def _inversion_geom_dependence_task(i, rm, pump_system, (depop_model_class1, depop_model_class2)):
    output.show_status((i, None), params.extended_status_strides, False)

    orig_geom = _set_geom(rm, params.beam_radius)
    try:
        active_medium = core.create_medium(None)

        depop_model1 = core.create_depop_model(active_medium, depop_model_class1)
        depop_model2 = core.create_depop_model(active_medium, depop_model_class2)

        inv1 = params.inverter_class(active_medium, pump_system, depop_model1)
        inv2 = params.inverter_class(active_medium, pump_system, depop_model2)

        inversion1 = inv1.invert(params.inversion_rtol, params.inversion_min_count_t)
        inversion2 = inv2.invert(params.inversion_rtol, params.inversion_min_count_t)

        gain_coef1 = inversion1 * active_medium.doping_agent.xsection
        gain_coef2 = inversion2 * active_medium.doping_agent.xsection

        gain1 = math.exp(gain_coef1 * active_medium.length)
        gain2 = math.exp(gain_coef2 * active_medium.length)

        stored_energy1 = model.energy.energy(params.lasing_wavelen, inversion1 * active_medium.volume)
        stored_energy2 = model.energy.energy(params.lasing_wavelen, inversion2 * active_medium.volume)

        inversion_rdiff = abs((inversion1 - inversion2) / min(inversion1, inversion2))
        gain_rdiff = abs((gain1 - gain2) / min(gain1, gain2))
    finally:
        _set_geom(*orig_geom)

    return inversion1, inversion2, gain1, gain2, stored_energy1, stored_energy2, inversion_rdiff, gain_rdiff

def compute_inversion_geom_dependence(task_pool, dirname):
    filename = lambda name: os.path.join(dirname, name)

    print output.div_line
    print "computing inversion dependence on geometry parameters"

    pump_system = model.pump.PumpSystem(params.pump_wavelen, params.pump_duration, params.pump_power, params.pump_efficiency)

    min_medium_radius = params.ext_opt_geom_mediumradius[0]

    count_rm = params.ext_opt_geom_resolution[0]

    Rm = np.linspace(min_medium_radius, params.ext_opt_geom_mediumradius[1], count_rm)

    depop_model_class1 = params.depop_model_class
    depop_model_class2 = params.ext_alt_depop_model

    inversions1, inversions2, gains1, gains2, stored_energies1, stored_energies2, inversion_rdiffs, gain_rdiffs = task_pool.parallel_task(_inversion_geom_dependence_task, (Rm,), (), (pump_system, (depop_model_class1, depop_model_class2)))

    output.show_status((count_rm, None), params.extended_status_strides, True)

    if params.graphs:
        print output.status_writing
        dirname = os.path.join(dirname, output.opt_geom_rel_path)
        dirname = output.init_dir(dirname)
        inversion_rdiff_max = params.ext_opt_inversion_rdiff_max
        pump_energy = params.pump_duration * params.pump_power
        energy_ylim = (0.0, pump_energy * 1.25)
        labels = [cls.descr for cls in [depop_model_class1, depop_model_class2]]
        plot.plot_data(filename("inversion"), "Inversion (%s)" % depop_model_class1.descr, (Rm, None, None, output.medium_radius_label), (inversions1, None, None, output.inversion_abs_label))
        plot.plot_data(filename("inversion_alt"), "Inversion (%s)" % depop_model_class2.descr, (Rm, None, None, output.medium_radius_label), (inversions2, None, None, output.inversion_abs_label))
        plot.plot_data(filename("inversions"), "Population Inversion", ([Rm]*2, None, None, output.medium_radius_label), ([inversions1, inversions2], None, None, output.inversion_abs_label), legend=labels)
        plot.plot_data(filename("ss_gain"), "Small Signal Gain (%s)" % depop_model_class1.descr, (Rm, None, None, output.medium_radius_label), (gains1, None, None, output.gain_label))
        plot.plot_data(filename("ss_gain_alt"), "Small Signal Gain (%s)" % depop_model_class2.descr, (Rm, None, None, output.medium_radius_label), (gains2, None, None, output.gain_label))
        plot.plot_data(filename("ss_gains"), "Small Signal Gain", ([Rm]*2, None, None, output.medium_radius_label), ([gains1, gains2], None, None, output.gain_label), legend=labels)
        plot.plot_data(filename("energy_stored"), "Stored Energy (%s)" % depop_model_class1.descr, (Rm, None, None, output.medium_radius_label), (stored_energies1, None, energy_ylim, output.energy_abs_stored_label), yvals=[(pump_energy, "pump energy")])
        plot.plot_data(filename("energy_stored_alt"), "Stored Energy (%s)" % depop_model_class2.descr, (Rm, None, None, output.medium_radius_label), (stored_energies2, None, energy_ylim, output.energy_abs_stored_label), yvals=[(pump_energy, "pump energy")])
        plot.plot_data(filename("energies_stored"), "Stored Energy", ([Rm]*2, None, None, output.medium_radius_label), ([stored_energies1, stored_energies2], None, energy_ylim, output.energy_abs_stored_label), legend=labels, yvals=[(pump_energy, "pump energy")])
        plot.plot_data(filename("inversion_rdiff"), "Inversion Relative Difference", (Rm, None, None, output.medium_radius_label), (inversion_rdiffs, None, None, output.inversion_rdiff_label), yvals=[(inversion_rdiff_max, None)])
        plot.plot_data(filename("ss_gain_rdiff"), "Small Signal Gain Relative Difference", (Rm, None, None, output.medium_radius_label), (gain_rdiffs, None, None, output.gain_rdiff_label), yvals=[(inversion_rdiff_max, None)])

    return inversions1, inversion_rdiffs

def _fluence_pump_dependence_task((i, j), (tau, pwr), inversion, active_medium, (rho, phi), (int_type, amp_type), (count_z, count_t), ref_pulse, lower_decay):
    output.show_status((i, j), params.extended_status_strides, False)

    initial_inversion = model.inversion.UniformInversion(inversion)
    active_medium.initial_inversion = initial_inversion

    integrator = model.integrator.DomainIntegrator(int_type)
    amp = amp_type(active_medium, count_z)

    fluence_out = _ref_signal_fluence(active_medium, (rho, phi), (integrator, amp), count_t, ref_pulse, lower_decay)
    return fluence_out

def compute_fluence_pump_dependence(task_pool, dirname, inversions, (int_type, amp_type), (count_z, count_t)):
    filename = lambda name: os.path.join(dirname, name)

    print output.div_line
    print "computing fluence dependence on pumping parameters"

    active_medium = core.create_medium(None)
    input_beam = core.create_beam()
    ref_pulse = core.create_pulse(active_medium, input_beam, input_beam.rho_ref, input_beam.phi_ref)
    pulse_train = core.create_train(ref_pulse)

    lower_decay = model.amplifier.lower_state_decay(active_medium, pulse_train)

    count_tau = params.ext_opt_pump_resolution[0]
    count_pwr = params.ext_opt_pump_resolution[1]

    Tau = np.linspace(params.ext_opt_pump_duration[0], params.ext_opt_pump_duration[1], count_tau)
    Pwr = np.linspace(params.ext_opt_pump_power[0], params.ext_opt_pump_power[1], count_pwr)

    rho, phi = input_beam.rho_ref, input_beam.phi_ref
    fluences = task_pool.parallel_task(_fluence_pump_dependence_task, (Tau, Pwr), (inversions,), (active_medium, (rho, phi), (int_type, amp_type), (count_z, count_t), ref_pulse, lower_decay))

    output.show_status((count_tau, count_pwr), params.extended_status_strides, True)

    if params.graphs:
        print output.status_writing
        dirname = os.path.join(dirname, output.opt_pump_rel_path)
        fluence_max = params.ext_opt_fluence_max
        zlim = None #(0.0, fluence_max)
        contours = [fluence_max]
        graph_types = [
            (dirname, Pwr, output.pump_power_label),
            (os.path.join(dirname, output.alt_plot_rel_path), Pwr * params.pump_efficiency / active_medium.volume, output.eff_power_density_label),
        ]
        for dirname, Y, ylabel in graph_types:
            dirname = output.init_dir(dirname)
            plot.plot_color(filename("fluence_out_max"), "Maximum Output Fluence", (Tau, None, None, output.pump_duration_label), (Y, None, None, ylabel), (fluences.T, zlim, output.fluence_abs_label_energy), params.out_num_auto_contours, contours)

    return fluences

def _fluence_geom_dependence_task((i, j), (rm, rb), inversion, (int_type, amp_type), (count_z, count_t)):
    output.show_status((i, j), params.extended_status_strides, False)

    orig_geom = _set_geom(rm, rb)
    try:
        active_medium = core.create_medium(inversion)
        input_beam = core.create_beam()
        rho, phi = input_beam.rho_ref, input_beam.phi_ref
        ref_pulse = core.create_pulse(active_medium, input_beam, rho, phi)
        pulse_train = core.create_train(ref_pulse)

        lower_decay = model.amplifier.lower_state_decay(active_medium, pulse_train)

        integrator = model.integrator.DomainIntegrator(int_type)
        amp = amp_type(active_medium, count_z)

        fluence_out = _ref_signal_fluence(active_medium, (rho, phi), (integrator, amp), count_t, ref_pulse, lower_decay)
    finally:
        _set_geom(*orig_geom)

    return fluence_out

def compute_fluence_geom_dependence(task_pool, dirname, inversions, (int_type, amp_type), (count_z, count_t)):
    filename = lambda name: os.path.join(dirname, name)

    print output.div_line
    print "computing fluence dependence on geometry parameters"

    min_medium_radius = params.ext_opt_geom_mediumradius[0]
    min_beam_radius = params.ext_opt_geom_beamradius[0]

    count_rm = params.ext_opt_geom_resolution[0]
    count_rb = params.ext_opt_geom_resolution[1]

    Rm = np.linspace(min_medium_radius, params.ext_opt_geom_mediumradius[1], count_rm)
    Rb = np.linspace(min_beam_radius, params.ext_opt_geom_beamradius[1], count_rb)

    inversions = np.meshgrid(inversions, Rb)[0].T
    fluences = task_pool.parallel_task(_fluence_geom_dependence_task, (Rm, Rb), (inversions,), ((int_type, amp_type), (count_z, count_t)))

    output.show_status((count_rm, count_rb), params.extended_status_strides, True)

    if params.graphs:
        print output.status_writing
        dirname = os.path.join(dirname, output.opt_geom_rel_path)
        dirname = output.init_dir(dirname)
        fluence_max = params.ext_opt_fluence_max
        zlim = None #(0.0, fluence_max)
        contours = [fluence_max]
        plot.plot_color(filename("fluence_out_max"), "Maximum Output Fluence", (Rm, None, None, output.medium_radius_label), (Rb, None, None, output.beam_radius_label), (fluences.T, zlim, output.fluence_abs_label_energy), params.out_num_auto_contours, contours)

    return fluences

def compute_pump_constraints(dirname, inversion_rdiffs, max_fluences):
    filename = lambda name: os.path.join(dirname, name)

    print output.div_line
    print "computing pumping parameters domain constraints"

    active_medium = core.create_medium(None)

    count_tau = params.ext_opt_pump_resolution[0]
    count_pwr = params.ext_opt_pump_resolution[1]

    Tau = np.linspace(params.ext_opt_pump_duration[0], params.ext_opt_pump_duration[1], count_tau)
    Pwr = np.linspace(params.ext_opt_pump_power[0], params.ext_opt_pump_power[1], count_pwr)

    inversion_rdiff_max = params.ext_opt_inversion_rdiff_max
    fluence_max = params.ext_opt_fluence_max
    contours = [
        (inversion_rdiffs.T, inversion_rdiff_max, "depopulation"),
        (max_fluences.T, fluence_max, "damage"),
    ]
    contour_comps = [1.0, 1.0]

    if params.graphs:
        print output.status_writing
        dirname = os.path.join(dirname, output.opt_pump_rel_path)
        graph_types = [
            (dirname, Pwr,output. pump_power_label),
            (os.path.join(dirname, output.alt_plot_rel_path), Pwr * params.pump_efficiency / active_medium.volume, output.eff_power_density_label),
        ]
        for dirname, Y, ylabel in graph_types:
            dirname = output.init_dir(dirname)
            plot.plot_contour(filename("constraints"), "Pumping Parameters Domain Constraints", (Tau, None, None, output.pump_duration_label), (Y, None, None, ylabel), contours)

    return (contours, None, None), (contour_comps, None, None)

def compute_geom_constraints(dirname, inversion_rdiffs, max_fluences):
    filename = lambda name: os.path.join(dirname, name)

    print output.div_line
    print "computing geometry parameters domain constraints"

    count_rm = params.ext_opt_geom_resolution[0]
    count_rb = params.ext_opt_geom_resolution[1]

    Rm = np.linspace(params.ext_opt_geom_mediumradius[0], params.ext_opt_geom_mediumradius[1], count_rm)
    Rb = np.linspace(params.ext_opt_geom_beamradius[0], params.ext_opt_geom_beamradius[1], count_rb)

    inversion_rdiff_max = params.ext_opt_inversion_rdiff_max
    fluence_max = params.ext_opt_fluence_max
    contours = [
        (max_fluences.T, fluence_max, "damage"),
    ]
    contour_comps = [1.0]

    rm_depop = np.interp(inversion_rdiff_max, inversion_rdiffs[::-1], Rm[::-1])
    xvals = [(rm_depop, "depopulation")]
    xval_comps = [-1.0]

    if params.graphs:
        print output.status_writing
        dirname = os.path.join(dirname, output.opt_geom_rel_path)
        dirname = output.init_dir(dirname)
        plot.plot_contour(filename("constraints"), "Geometry Parameters Domain Constraints", (Rm, None, None, output.medium_radius_label), (Rb, None, None, output.beam_radius_label), contours, xvals=xvals)

    return (contours, xvals, None), (contour_comps, xval_comps, None)

def optimize_output(domain, outputs, limits, comparisons, price):
    params1, params2 = domain
    contours, xvals, yvals = limits
    contour_comps, xval_comps, yval_comps = comparisons
    optimum = None
    for i, p1 in enumerate(params1):
        for j, p2 in enumerate(params2):
            safe = True
            if contours:
                for limit, comp in zip(contours, contour_comps):
                    Z, zmax, _ = limit
                    Z = Z.T
                    if Z[i, j] * comp > zmax * comp:
                        safe = False
                        break
            if xvals:
                for limit, comp in zip(xvals, xval_comps):
                    val, _ = limit
                    if p1 * comp > val * comp:
                        safe = False
                        break
            if yvals:
                for limit, comp in zip(yvals, yval_comps):
                    val, _ = limit
                    if p2 *comp > val * comp:
                        safe = False
                        break
            if not safe:
                continue
            if not optimum:
                optimum = (i, j)
            else:
                output = outputs[i, j]
                optimum_output = outputs[optimum]
                if output > optimum_output:
                    optimum = (i, j)
                elif output == optimum_output:
                    current_price = price(p1, p2)
                    optimum_price = price(params1[optimum[0]], params2[optimum[1]])
                    if current_price < optimum_price:
                        optimum = (i, j)
    return optimum

def _energy_pump_dependence_task((i, j), (tau, pwr), inversion, num_types, counts):
    output.show_status((i, j), params.extended_status_strides, False)
    _, _, output_energy, rel_gain_decrease = core.amplify_train(None, num_types, counts, inversion, quiet=True)
    return output_energy, rel_gain_decrease

def compute_energy_pump_dependence(task_pool, dirname, inversions, constraints, num_types, counts):
    filename = lambda name: os.path.join(dirname, name)

    print output.div_line
    print "computing energy dependence on pumping parameters"

    active_medium = core.create_medium(None)
    input_beam = core.create_beam()

    input_photon_count = input_beam.fluence_integral(active_medium.radius)
    input_energy = model.energy.energy(params.lasing_wavelen, input_photon_count)
    input_energy *= params.train_pulse_count

    count_tau = params.ext_opt_pump_resolution[0]
    count_pwr = params.ext_opt_pump_resolution[1]

    Tau = np.linspace(params.ext_opt_pump_duration[0], params.ext_opt_pump_duration[1], count_tau)
    Pwr = np.linspace(params.ext_opt_pump_power[0], params.ext_opt_pump_power[1], count_pwr)

    output_energies, rel_gain_decreases = task_pool.parallel_task(_energy_pump_dependence_task, (Tau, Pwr), (inversions,), (num_types, counts))

    output.show_status((count_tau, count_pwr), params.extended_status_strides, True)

    pump_energies = np.prod(np.array(np.meshgrid(Tau, Pwr)), axis=0).T
    stored_energies = model.energy.energy(params.lasing_wavelen, inversions * active_medium.volume)
    energy_gains = output_energies / input_energy
    added_energies = output_energies - input_energy
    extraction_effs = added_energies / stored_energies
    total_effs = added_energies / pump_energies

    limits, comparisons = constraints
    price = lambda tau, pwr: tau * pwr

    unitconv.print_result("input energy [{}]: {}", ("mJ",), (input_energy,))

    optimum = optimize_output((Tau, Pwr), output_energies, limits, comparisons, price)
    output_energy_optimum_params = (Tau[optimum[0]], Pwr[optimum[1]]) if optimum else (None, None)
    unitconv.print_result("max. output energy [{}]: {}", ("mJ",), (output_energies[optimum] if optimum else None,))
    unitconv.print_result("optimum pumping parameters (duration [{}], power [{}]): ({}, {})", ("us", "W"), output_energy_optimum_params)

    optimum = optimize_output((Tau, Pwr), energy_gains, limits, comparisons, price)
    energy_gain_optimum_params = (Tau[optimum[0]], Pwr[optimum[1]]) if optimum else (None, None)
    unitconv.print_result("max. energy gain: {}", (), (energy_gains[optimum] if optimum else None,))
    unitconv.print_result("optimum pumping parameters (duration [{}], power [{}]): ({}, {})", ("us", "W"), energy_gain_optimum_params)

    optimum = optimize_output((Tau, Pwr), extraction_effs, limits, comparisons, price)
    extraction_eff_optimum_params = (Tau[optimum[0]], Pwr[optimum[1]]) if optimum else (None, None)
    unitconv.print_result("max. extraction efficiency [{}]: {}", ("%",), (extraction_effs[optimum] if optimum else None,))
    unitconv.print_result("optimum pumping parameters (duration [{}], power [{}]): ({}, {})", ("us", "W"), extraction_eff_optimum_params)

    optimum = optimize_output((Tau, Pwr), total_effs, limits, comparisons, price)
    total_eff_optimum_params = (Tau[optimum[0]], Pwr[optimum[1]]) if optimum else (None, None)
    unitconv.print_result("max. opt.-opt. efficiency [{}]: {}", ("%",), (total_effs[optimum] if optimum else None,))
    unitconv.print_result("optimum pumping parameters (duration [{}], power [{}]): ({}, {})", ("us", "W"), total_eff_optimum_params)

    optimum = optimize_output((Tau, Pwr), -rel_gain_decreases, limits, comparisons, price)
    rel_gain_decrease_optimum_params = (Tau[optimum[0]], Pwr[optimum[1]]) if optimum else (None, None)
    unitconv.print_result("min. rel. gain decrease [{}]: {}", ("%",), (rel_gain_decreases[optimum] if optimum else None,))
    unitconv.print_result("optimum pumping parameters (duration [{}], power [{}]): ({}, {})", ("us", "W"), rel_gain_decrease_optimum_params)

    if params.graphs:
        print output.status_writing
        extra_contours, xvals, yvals = limits
        dirname = os.path.join(dirname, output.opt_pump_rel_path)
        graph_types = [
            (dirname, Pwr, output.pump_power_label),
            (os.path.join(dirname, output.alt_plot_rel_path), Pwr * params.pump_efficiency / active_medium.volume, output.eff_power_density_label),
        ]
        for dirname, Y, ylabel in graph_types:
            dirname = output.init_dir(dirname)
            plot.plot_color(filename("energy_out"), "Output Energy", (Tau, None, None, output.pump_duration_label), (Y, None, None, ylabel), (output_energies.T, None, output.energy_abs_pulse_label), params.out_num_auto_contours, extra_contours=extra_contours, xvals=xvals, yvals=yvals)
            plot.plot_color(filename("energy_gain"), "Energy Gain", (Tau, None, None, output.pump_duration_label), (Y, None, None, ylabel), (energy_gains.T, None, output.energy_rel_label), params.out_num_auto_contours, extra_contours=extra_contours, xvals=xvals, yvals=yvals)
            plot.plot_color(filename("efficiency_extr"), "Extraction Efficiency", (Tau, None, None, output.pump_duration_label), (Y, None, None, ylabel), (extraction_effs.T, None, output.extraction_eff_label), params.out_num_auto_contours, extra_contours=extra_contours, xvals=xvals, yvals=yvals)
            plot.plot_color(filename("efficiency_opt2"), "Optical to Optical Efficiency", (Tau, None, None, output.pump_duration_label), (Y, None, None, ylabel), (total_effs.T, None, output.total_eff_label), params.out_num_auto_contours, extra_contours=extra_contours, xvals=xvals, yvals=yvals)
            if params.train_pulse_count > 1:
                plot.plot_color(filename("gain_decrease"), "Gain Decrease", (Tau, None, None, output.pump_duration_label), (Y, None, None, ylabel), (rel_gain_decreases.T, None, output.rel_gain_decrease_label), params.out_num_auto_contours, extra_contours=extra_contours, xvals=xvals, yvals=yvals)

def _energy_geom_dependence_task((i, j), (rm, rb), inversion, num_types, counts):
    output.show_status((i, j), params.extended_status_strides, False)

    orig_geom = _set_geom(rm, rb)
    try:
        active_medium = core.create_medium(None)
        input_beam = core.create_beam()
        stored_energy = model.energy.energy(params.lasing_wavelen, inversion * active_medium.volume)
        input_photon_count = input_beam.fluence_integral(active_medium.radius)
        input_energy = model.energy.energy(params.lasing_wavelen, input_photon_count)
        input_energy *= params.train_pulse_count

        _, _, output_energy, rel_gain_decrease = core.amplify_train(None, num_types, counts, inversion, quiet=True)
    finally:
        _set_geom(*orig_geom)

    return stored_energy, input_energy, output_energy, rel_gain_decrease

def compute_energy_geom_dependence(task_pool, dirname, inversions, constraints, num_types, counts):
    filename = lambda name: os.path.join(dirname, name)

    print output.div_line
    print "computing energy dependence on geometry parameters"

    min_medium_radius = params.ext_opt_geom_mediumradius[0]
    min_beam_radius = params.ext_opt_geom_beamradius[0]

    count_rm = params.ext_opt_geom_resolution[0]
    count_rb = params.ext_opt_geom_resolution[1]

    Rm = np.linspace(min_medium_radius, params.ext_opt_geom_mediumradius[1], count_rm)
    Rb = np.linspace(min_beam_radius, params.ext_opt_geom_beamradius[1], count_rb)

    inversions = np.meshgrid(inversions, Rb)[0].T
    stored_energies, input_energies, output_energies, rel_gain_decreases = task_pool.parallel_task(_energy_geom_dependence_task, (Rm, Rb), (inversions,), (num_types, counts))

    output.show_status((count_rm, count_rb), params.extended_status_strides, True)

    pump_energy = params.pump_duration * params.pump_power
    energy_gains = output_energies / input_energies
    added_energies = output_energies - input_energies
    extraction_effs = added_energies / stored_energies
    total_effs = added_energies / pump_energy

    limits, comparisons = constraints
    price = lambda rm, rb: 1.0 / (rm * rb)

    optimum = optimize_output((Rm, Rb), output_energies, limits, comparisons, price)
    output_energy_optimum_params = (Rm[optimum[0]], Rb[optimum[1]]) if optimum else (None, None)
    unitconv.print_result("max. output energy [{}]: {}", ("mJ",), (output_energies[optimum] if optimum else None,))
    unitconv.print_result("optimum geometry parameters (medium diameter [{}], beam diameter [{}]): ({}, {})", ("mm", "mm"), output_energy_optimum_params)

    optimum = optimize_output((Rm, Rb), energy_gains, limits, comparisons, price)
    energy_gain_optimum_params = (Rm[optimum[0]], Rb[optimum[1]]) if optimum else (None, None)
    unitconv.print_result("max. energy gain: {}", (), (energy_gains[optimum] if optimum else None,))
    unitconv.print_result("optimum geometry parameters (medium diameter [{}], beam diameter [{}]): ({}, {})", ("mm", "mm"), energy_gain_optimum_params)

    optimum = optimize_output((Rm, Rb), extraction_effs, limits, comparisons, price)
    extraction_eff_optimum_params = (Rm[optimum[0]], Rb[optimum[1]]) if optimum else (None, None)
    unitconv.print_result("max. extraction efficiency [{}]: {}", ("%",), (extraction_effs[optimum] if optimum else None,))
    unitconv.print_result("optimum geometry parameters (medium diameter [{}], beam diameter [{}]): ({}, {})", ("mm", "mm"), extraction_eff_optimum_params)

    optimum = optimize_output((Rm, Rb), total_effs, limits, comparisons, price)
    total_eff_optimum_params = (Rm[optimum[0]], Rb[optimum[1]]) if optimum else (None, None)
    unitconv.print_result("max. opt.-opt. efficiency [{}]: {}", ("%",), (total_effs[optimum] if optimum else None,))
    unitconv.print_result("optimum geometry parameters (medium diameter [{}], beam diameter [{}]): ({}, {})", ("mm", "mm"), total_eff_optimum_params)

    optimum = optimize_output((Rm, Rb), -rel_gain_decreases, limits, comparisons, price)
    rel_gain_decrease_optimum_params = (Rm[optimum[0]], Rb[optimum[1]]) if optimum else (None, None)
    unitconv.print_result("min. rel. gain decrease [{}]: {}", ("%",), (rel_gain_decreases[optimum] if optimum else None,))
    unitconv.print_result("optimum geometry parameters (medium diameter [{}], beam diameter [{}]): ({}, {})", ("mm", "mm"), rel_gain_decrease_optimum_params)

    if params.graphs:
        print output.status_writing
        extra_contours, xvals, yvals = limits
        dirname = os.path.join(dirname, output.opt_geom_rel_path)
        dirname = output.init_dir(dirname)
        plot.plot_color(filename("energy_in"), "Input Energy", (Rm, None, None, output.medium_radius_label), (Rb, None, None, output.beam_radius_label), (input_energies.T, None, output.energy_abs_pulse_label), params.out_num_auto_contours, extra_contours=extra_contours, xvals=xvals, yvals=yvals)
        plot.plot_color(filename("energy_out"), "Output Energy", (Rm, None, None, output.medium_radius_label), (Rb, None, None, output.beam_radius_label), (output_energies.T, None, output.energy_abs_pulse_label), params.out_num_auto_contours, extra_contours=extra_contours, xvals=xvals, yvals=yvals)
        plot.plot_color(filename("energy_gain"), "Energy Gain", (Rm, None, None, output.medium_radius_label), (Rb, None, None, output.beam_radius_label), (energy_gains.T, None, output.energy_rel_label), params.out_num_auto_contours, extra_contours=extra_contours, xvals=xvals, yvals=yvals)
        plot.plot_color(filename("efficiency_extr"), "Extraction Efficiency", (Rm, None, None, output.medium_radius_label), (Rb, None, None, output.beam_radius_label), (extraction_effs.T, None, output.extraction_eff_label), params.out_num_auto_contours, extra_contours=extra_contours, xvals=xvals, yvals=yvals)
        plot.plot_color(filename("efficiency_opt2"), "Optical to Optical Efficiency", (Rm, None, None, output.medium_radius_label), (Rb, None, None, output.beam_radius_label), (total_effs.T, None, output.total_eff_label), params.out_num_auto_contours, extra_contours=extra_contours, xvals=xvals, yvals=yvals)
        if params.train_pulse_count > 1:
            plot.plot_color(filename("gain_decrease"), "Gain Decrease", (Rm, None, None, output.medium_radius_label), (Rb, None, None, output.beam_radius_label), (rel_gain_decreases.T, None, output.rel_gain_decrease_label), params.out_num_auto_contours, extra_contours=extra_contours, xvals=xvals, yvals=yvals)

def compare_lower_lifetimes(dirname, ref_inversion, (int_types, amp_types), numerics):
    filename = lambda name: os.path.join(dirname, name)

    print output.div_line
    print "comparing lower state lifetimes"

    if numerics is None:
        numerics, _ = core.select_methods((int_types, amp_types), ref_inversion, quiet=True)
    num_types, counts = numerics

    lower_lifetime = params.dopant_lower_lifetime

    active_medium = core.create_medium(ref_inversion)
    active_medium_3 = copy.deepcopy(active_medium)
    active_medium_3.doping_agent.lower_lifetime = float("inf")
    active_medium_4 = copy.deepcopy(active_medium)
    active_medium_4.doping_agent.lower_lifetime = 0.0

    input_beam = core.create_beam()

    ref_pulse = core.create_pulse(active_medium, input_beam, input_beam.rho_ref, input_beam.phi_ref)

    (int_type, amp_type), (_, _, count_z, count_t) = num_types, counts

    integrator = model.integrator.DomainIntegrator(int_type)

    amp = amp_type(active_medium, count_z)
    amp_3 = model.amplifier.ExactOutputAmplifier(active_medium_3, count_z)
    amp_4 = model.amplifier.ExactOutputAmplifier(active_medium_4, count_z)
    amplify_args = (input_beam.rho_ref, input_beam.phi_ref, ref_pulse, count_t)

    print "zero"
    density_out_4, _ = amp_4.amplify(*amplify_args)
    fluence_out_4 = integrator.integrate(amp_4.T, density_out_4) * active_medium_4.light_speed
    fluence_gain_4 = fluence_out_4 / input_beam.ref_fluence
    unitconv.print_result("fluence gain: {}", (), (fluence_gain_4,))

    lsl_output_label = "finite"
    if lower_lifetime == 0.0:
        lsl_output_label = "zero"
    elif math.isinf(lower_lifetime):
        lsl_output_label = "infinite"
    print lsl_output_label
    density_out, _ = amp.amplify(*amplify_args)
    fluence_out = integrator.integrate(amp.T, density_out) * active_medium.light_speed
    fluence_gain = fluence_out / input_beam.ref_fluence
    unitconv.print_result("fluence gain: {}", (), (fluence_gain,))

    print "infinite"
    density_out_3, _ = amp_3.amplify(*amplify_args)
    fluence_out_3 = integrator.integrate(amp_3.T, density_out_3) * active_medium_3.light_speed
    fluence_gain_3 = fluence_out_3 / input_beam.ref_fluence
    unitconv.print_result("fluence gain: {}", (), (fluence_gain_3,))

    if params.graphs:
        print output.status_writing
        dirname = os.path.join(dirname, output.ref_pulse_rel_path)
        dirname = output.init_dir(dirname)
        T = amp.T
        if params.output_rel_time:
            T = T / params.pulse_duration
        out_t_label = output.norm_t_label if params.output_rel_time else output.t_amp_label
        tlim = (T[0], T[-1])
        Ts = (T,) * 3
        ref_density = ref_pulse.ref_density
        densities = (density_out_4 / ref_density, density_out / ref_density, density_out_3 / ref_density)
        lifetime_scale = unitconv.units[output.lower_lifetime_unit]
        lsl_graph_label_fmt = (r"%g \, %s" % (lower_lifetime/lifetime_scale, output.lower_lifetime_unit))
        if lower_lifetime == 0.0:
            lsl_graph_label_fmt = "0"
        elif math.isinf(lower_lifetime):
            lsl_graph_label_fmt = "\\infty"
        labels = (output.lower_lifetime_legend % "0", output.lower_lifetime_legend % lsl_graph_label_fmt, output.lower_lifetime_legend % "\\infty")
        plot.plot_data(filename("lsl_effects"), "Effects of Lower State Lifetime", (Ts, None, tlim, out_t_label), (densities, None, None, output.density_rel_label), labels)

def select_methods(perform_opt_pump, perform_opt_geom, (int_types, amp_types), inversions_pump, inversions_geom):
    print output.div_line
    print "determining extended mode method combinations"

    (num_types_pump, counts_pump), (num_types_geom, counts_geom) = (None, None), (None, None)

    if perform_opt_pump:
        print "pumping"
        max_inversion_pump = inversions_pump[-1, -1]
        (num_types_pump, counts_pump), _ = core.select_methods((int_types, amp_types), max_inversion_pump, quiet=True)

    if perform_opt_geom:
        print "geometry"
        max_medium_radius = params.ext_opt_geom_mediumradius[1]
        min_beam_radius = params.ext_opt_geom_beamradius[0]

        orig_geom = _set_geom(max_medium_radius, min_beam_radius)
        try:
            max_inversion_geom = inversions_geom[0]
            (num_types_geom, counts_geom), _ = core.select_methods((int_types, amp_types), max_inversion_geom, quiet=True)
        finally:
            _set_geom(*orig_geom)

    return (num_types_pump, counts_pump), (num_types_geom, counts_geom)

def extended_mode(task_pool, dirname, ref_inversion, (int_types, amp_types), numerics):
    if params.amplification:
        compare_lower_lifetimes(dirname, ref_inversion, (int_types, amp_types), numerics)

    if not params.initial_inversion:
        compare_depop_models(dirname)

        perform_opt_pump = min(params.ext_opt_pump_resolution) > 1
        perform_opt_geom = min(params.ext_opt_geom_resolution) > 1
        perform_opt = perform_opt_pump or perform_opt_geom

        if perform_opt:
            inversions_pump, inversions_geom = None, None
            if perform_opt_pump:
                inversions_pump, inversion_rdiffs_pump = compute_inversion_pump_dependence(task_pool, dirname)
            if perform_opt_geom:
                inversions_geom, inversion_rdiffs_geom = compute_inversion_geom_dependence(task_pool, dirname)

            if params.amplification:
                (num_types_pump, counts_pump), (num_types_geom, counts_geom) = select_methods(perform_opt_pump, perform_opt_geom, (int_types, amp_types), inversions_pump, inversions_geom)

                if perform_opt_pump:
                    _, _, count_z_pump, count_t_pump = counts_pump
                    max_fluences_pump = compute_fluence_pump_dependence(task_pool, dirname, inversions_pump, num_types_pump, (count_z_pump, count_t_pump))
                if perform_opt_geom:
                    _, _, count_z_geom, count_t_geom = counts_geom
                    max_fluences_geom = compute_fluence_geom_dependence(task_pool, dirname, inversions_geom, num_types_geom, (count_z_geom, count_t_geom))

                if perform_opt_pump:
                    pump_constraints = compute_pump_constraints(dirname, inversion_rdiffs_pump, max_fluences_pump)
                if perform_opt_geom:
                    geom_constraints = compute_geom_constraints(dirname, inversion_rdiffs_geom, max_fluences_geom)

                if perform_opt_pump:
                    compute_energy_pump_dependence(task_pool, dirname, inversions_pump, pump_constraints, num_types_pump, counts_pump)
                if perform_opt_geom:
                    compute_energy_geom_dependence(task_pool, dirname, inversions_geom, geom_constraints, num_types_geom, counts_geom)


#!/usr/bin/env python
# Copyright 2014-2021 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

'''
spin-free X2C correction for extended systems (experimental feature)
'''

from functools import reduce
import copy
import numpy
import scipy.linalg
from pyscf import lib
from pyscf.gto import mole
from pyscf.lib import logger
from pyscf.x2c import x2c
from pyscf.pbc import gto as pbcgto
from pyscf.pbc import tools
from pyscf.pbc.df import aft
from pyscf.pbc.df import aft_jk
from pyscf.pbc.df import ft_ao
from pyscf.pbc.scf import ghf
from pyscf import __config__

def sfx2c1e(mf):
    '''Spin-free X2C.
    For the given SCF object, update the hcore constructor.

    Args:
        mf : an SCF object

    Returns:
        An SCF object

    Examples:

    >>> mol = gto.M(atom='H 0 0 0; F 0 0 1', basis='ccpvdz', verbose=0)
    >>> mf = scf.RHF(mol).sfx2c1e()
    >>> mf.scf()

    >>> mol.symmetry = 1
    >>> mol.build(0, 0)
    >>> mf = scf.UHF(mol).sfx2c1e()
    >>> mf.scf()
    '''
    if isinstance(mf, x2c._X2C_SCF):
        if mf.with_x2c is None:
            return mf.__class__(mf)
        else:
            return mf

    mf_class = mf.__class__
    if mf_class.__doc__ is None:
        doc = ''
    else:
        doc = mf_class.__doc__
    class SFX2C1E_SCF(mf_class, x2c._X2C_SCF):
        __doc__ = doc + '''
        Attributes for spin-free X2C:
            with_x2c : X2C object
        '''
        def __init__(self, mf):
            self.__dict__.update(mf.__dict__)
            self.with_x2c = SpinFreeX2C(mf.mol)
            self._keys = self._keys.union(['with_x2c'])

        def get_hcore(self, cell=None, kpts=None, kpt=None):
            if cell is None: cell = self.cell
            if kpts is None:
                if getattr(self, 'kpts', None) is not None:
                    kpts = self.kpts
                else:
                    if kpt is None:
                        kpts = self.kpt
                    else:
                        kpts = kpt
            if self.with_x2c:
                hcore = self.with_x2c.get_hcore(cell, kpts)
                if isinstance(self, ghf.GHF):
                    if kpts.ndim == 1:
                        hcore = scipy.linalg.block_diag(hcore, hcore)
                    else:
                        hcore = [scipy.linalg.block_diag(h, h) for h in hcore]
                return hcore
            else:
                return mf_class.get_hcore(self, cell, kpts)

    return SFX2C1E_SCF(mf)

sfx2c = sfx2c1e

class X2C(x2c.X2C):

    exp_drop = getattr(__config__, 'pbc_x2c_X2C_exp_drop', 0.2)
    approx = getattr(__config__, 'pbc_x2c_X2C_approx', 'atom1e')
    xuncontract = getattr(__config__, 'pbc_x2c_X2C_xuncontract', True)
    basis = getattr(__config__, 'pbc_x2c_X2C_basis', None)

    def __init__(self, cell, kpts=None):
        self.cell = cell
        x2c.X2C.__init__(self, cell)

class SpinFreeX2C(X2C):
    def get_hcore(self, cell=None, kpts=None):
        if cell is None: cell = self.cell
        if kpts is None:
            kpts_lst = numpy.zeros((1,3))
        else:
            kpts_lst = numpy.reshape(kpts, (-1,3))

        xcell, contr_coeff = self.get_xmol(cell)
        with_df = aft.AFTDF(xcell)
        c = lib.param.LIGHT_SPEED
        assert('1E' in self.approx.upper())
        if 'ATOM' in self.approx.upper():
            atom_slices = xcell.offset_nr_by_atom()
            nao = xcell.nao_nr()
            x = numpy.zeros((nao,nao))
            vloc = numpy.zeros((nao,nao))
            wloc = numpy.zeros((nao,nao))
            for ia in range(xcell.natm):
                ish0, ish1, p0, p1 = atom_slices[ia]
                shls_slice = (ish0, ish1, ish0, ish1)
                t1 = xcell.intor('int1e_kin', shls_slice=shls_slice)
                s1 = xcell.intor('int1e_ovlp', shls_slice=shls_slice)
                with xcell.with_rinv_at_nucleus(ia):
                    z = -xcell.atom_charge(ia)
                    v1 = z * xcell.intor('int1e_rinv', shls_slice=shls_slice)
                    w1 = z * xcell.intor('int1e_prinvp', shls_slice=shls_slice)
                vloc[p0:p1,p0:p1] = v1
                wloc[p0:p1,p0:p1] = w1
                x[p0:p1,p0:p1] = x2c._x2c1e_xmatrix(t1, v1, w1, s1, c)
        else:
            raise NotImplementedError

        t = xcell.pbc_intor('int1e_kin', 1, lib.HERMITIAN, kpts_lst)
        s = xcell.pbc_intor('int1e_ovlp', 1, lib.HERMITIAN, kpts_lst)
        v = with_df.get_nuc(kpts_lst)
        #w = get_pnucp(with_df, kpts_lst)
        if self.basis is not None:
            s22 = s
            s21 = pbcgto.intor_cross('int1e_ovlp', xcell, cell, kpts=kpts_lst)

        h1_kpts = []
        for k in range(len(kpts_lst)):
            # The treatment of pnucp local part has huge effects to hcore
            #h1 = x2c._get_hcore_fw(t[k], vloc, wloc, s[k], x, c) - vloc + v[k]
            #h1 = x2c._get_hcore_fw(t[k], v[k], w[k], s[k], x, c)
            h1 = x2c._get_hcore_fw(t[k], v[k], wloc, s[k], x, c)
            if self.basis is not None:
                c = lib.cho_solve(s22[k], s21[k])
                h1 = reduce(numpy.dot, (c.T, h1, c))
            if self.xuncontract and contr_coeff is not None:
                h1 = reduce(numpy.dot, (contr_coeff.T, h1, contr_coeff))
            h1_kpts.append(h1)

        if kpts is None or numpy.shape(kpts) == (3,):
            h1_kpts = h1_kpts[0]
        return lib.asarray(h1_kpts)

    def get_xmat(self, cell=None, kpts=None):
        if cell is None: cell = self.cell
        xcell, contr_coeff = self.get_xmol(cell)
        c = lib.param.LIGHT_SPEED
        assert('1E' in self.approx.upper())
        if 'ATOM' in self.approx.upper():
            atom_slices = xcell.offset_nr_by_atom()
            nao = xcell.nao_nr()
            x = numpy.zeros((nao,nao))
            for ia in range(xcell.natm):
                ish0, ish1, p0, p1 = atom_slices[ia]
                shls_slice = (ish0, ish1, ish0, ish1)
                t1 = xcell.intor('int1e_kin', shls_slice=shls_slice)
                s1 = xcell.intor('int1e_ovlp', shls_slice=shls_slice)
                with xcell.with_rinv_at_nucleus(ia):
                    z = -xcell.atom_charge(ia)
                    v1 = z * xcell.intor('int1e_rinv', shls_slice=shls_slice)
                    w1 = z * xcell.intor('int1e_prinvp', shls_slice=shls_slice)
                x[p0:p1,p0:p1] = x2c._x2c1e_xmatrix(t1, v1, w1, s1, c)
        else:
            raise NotImplementedError
        return x

# Use Ewald-like technique to compute spVsp.
# spVsp may not be divergent because the numerator spsp and the denominator
# in Coulomb kernel 4pi/G^2 are likely cancelled.  Even a real space lattice
# sum can converge to a finite value, it's difficult to accurately converge
# this value, i.e., large number of images in lattice summation is required.
def get_pnucp(mydf, kpts=None):
    cell = mydf.cell
    if kpts is None:
        kpts_lst = numpy.zeros((1,3))
    else:
        kpts_lst = numpy.reshape(kpts, (-1,3))

    log = logger.Logger(mydf.stdout, mydf.verbose)
    t1 = (logger.process_clock(), logger.perf_counter())

    nkpts = len(kpts_lst)
    nao = cell.nao_nr()
    nao_pair = nao * (nao+1) // 2

    Gv, Gvbase, kws = cell.get_Gv_weights(mydf.mesh)
    charge = -cell.atom_charges()
    kpt_allow = numpy.zeros(3)
    coulG = tools.get_coulG(cell, kpt_allow, mesh=mydf.mesh, Gv=Gv)
    coulG *= kws
    if mydf.eta == 0:
        wj = numpy.zeros((nkpts,nao_pair), dtype=numpy.complex128)
        SI = cell.get_SI(Gv)
        vG = numpy.einsum('i,ix->x', charge, SI) * coulG
        wj = numpy.zeros((nkpts,nao_pair), dtype=numpy.complex128)

    else:
        nuccell = copy.copy(cell)
        half_sph_norm = .5/numpy.sqrt(numpy.pi)
        norm = half_sph_norm/mole.gaussian_int(2, mydf.eta)
        chg_env = [mydf.eta, norm]
        ptr_eta = cell._env.size
        ptr_norm = ptr_eta + 1
        chg_bas = [[ia, 0, 1, 1, 0, ptr_eta, ptr_norm, 0] for ia in range(cell.natm)]
        nuccell._atm = cell._atm
        nuccell._bas = numpy.asarray(chg_bas, dtype=numpy.int32)
        nuccell._env = numpy.hstack((cell._env, chg_env))

        wj = lib.asarray(mydf._int_nuc_vloc(nuccell, kpts_lst, 'int3c2e_pvp1'))
        t1 = log.timer_debug1('pnucp pass1: analytic int', *t1)

        aoaux = ft_ao.ft_ao(nuccell, Gv)
        vG = numpy.einsum('i,xi->x', charge, aoaux) * coulG
        if cell.dimension == 3:
            nucbar = sum([z/nuccell.bas_exp(i)[0] for i,z in enumerate(charge)])
            nucbar *= numpy.pi/cell.vol

            ovlp = cell.pbc_intor('int1e_kin', 1, lib.HERMITIAN, kpts_lst)
            for k in range(nkpts):
                s = lib.pack_tril(ovlp[k])
                # *2 due to the factor 1/2 in T
                wj[k] -= nucbar*2 * s

    max_memory = max(2000, mydf.max_memory-lib.current_memory()[0])
    for aoaoks, p0, p1 in mydf.ft_loop(mydf.mesh, kpt_allow, kpts_lst,
                                       max_memory=max_memory, aosym='s2',
                                       intor='GTO_ft_pdotp'):
        for k, aoao in enumerate(aoaoks):
            if aft_jk.gamma_point(kpts_lst[k]):
                wj[k] += numpy.einsum('k,kx->x', vG[p0:p1].real, aoao.real)
                wj[k] += numpy.einsum('k,kx->x', vG[p0:p1].imag, aoao.imag)
            else:
                wj[k] += numpy.einsum('k,kx->x', vG[p0:p1].conj(), aoao)
    t1 = log.timer_debug1('contracting pnucp', *t1)

    wj_kpts = []
    for k, kpt in enumerate(kpts_lst):
        if aft_jk.gamma_point(kpt):
            wj_kpts.append(lib.unpack_tril(wj[k].real.copy()))
        else:
            wj_kpts.append(lib.unpack_tril(wj[k]))

    if kpts is None or numpy.shape(kpts) == (3,):
        wj_kpts = wj_kpts[0]
    return numpy.asarray(wj_kpts)

if __name__ == '__main__':
    from pyscf.pbc import scf
    cell = pbcgto.Cell()
    cell.build(unit = 'B',
               a = numpy.eye(3)*4,
               mesh = [11]*3,
               atom = 'H 0 0 0; H 0 0 1.8',
               verbose = 4,
               basis='sto3g')
    lib.param.LIGHT_SPEED = 2
    mf = scf.RHF(cell)
    mf.with_df = aft.AFTDF(cell)
    enr = mf.kernel()
    print('E(NR) = %.12g' % enr)

    mf = sfx2c1e(mf)
    esfx2c = mf.kernel()
    print('E(SFX2C1E) = %.12g' % esfx2c)

    mf = scf.KRHF(cell)
    mf.with_df = aft.AFTDF(cell)
    mf.kpts = cell.make_kpts([2,2,1])
    enr = mf.kernel()
    print('E(k-NR) = %.12g' % enr)

    mf = sfx2c1e(mf)
    esfx2c = mf.kernel()
    print('E(k-SFX2C1E) = %.12g' % esfx2c)

#    cell = pbcgto.M(unit = 'B',
#               a = numpy.eye(3)*4,
#               atom = 'H 0 0 0; H 0 0 1.8',
#               mesh = None,
#               dimension = 2,
#               basis='sto3g')
#    with_df = aft.AFTDF(cell)
#    w0 = get_pnucp(with_df, cell.make_kpts([2,2,1]))
#    with_df = aft.AFTDF(cell)
#    with_df.eta = 0
#    w1 = get_pnucp(with_df, cell.make_kpts([2,2,1]))
#    print(abs(w0-w1).max())


#!/usr/bin/env python

import sys, time,ast
from copy import copy, deepcopy
import numpy as np
import argparse
import gc
import cPickle as pickle

class alchemy:
   def getpair(self, sa, sb):
      if len(self.rules)==0: # special case when the alchemical matrix is default
         if sa==sb: return 1
         else: return 0
      else:
          if sa<=sb and (sa,sb) in self.rules:
             return self.rules[(sa,sb)]
          elif sa>sb and (sb,sa) in self.rules:
             return self.rules[(sb,sa)]
          else:
             if sa==sb: return 1
             else: return 0

   def __init__(self, rules={}, mu=0):
      self.rules = rules.copy()
      self.mu = mu

class alchemy_mendeleev:
   def getpair(self, sa, sb):
      if sa<=sb and (sa,sb) in self.rules:
         return self.rules[(sa,sb)]
      elif sa>sb and (sb,sa) in self.rules:
         return self.rules[(sb,sa)]
      else:
         Elec_neg={1: 2.2, 2: 0, 3: 0.98, 4: 1.57, 5: 2.04, 6: 2.55, 7: 3.04, 8: 3.44, 9: 3.98, 10: 0, 11: 0.93, 12: 1.31, 13: 1.5, 14: 1.8, 15: 2.19, 16: 2.58, 17: 3.16, 18: 0, 19: 0.82, 20: 1, 21: 1.36, 22: 1.54, 23: 1.63, 24: 1.66, 25: 1.55, 26: 1.83, 27: 1.88, 28: 1.91, 29: 1.9, 30: 1.65, 31: 1.81, 32: 2.01, 33: 2.18, 34: 2.55, 35: 2.96, 36: 0, 37: 0.82, 38: 0.95, 39: 1.22, 40: 1.33, 41: 1.6, 42: 2.16, 43: 1.9, 44: 2.2, 45: 2.28, 46: 2.2, 47: 1.93, 48: 1.69, 49: 1.78, 50: 1.96, 51: 2.05, 52: 2.1, 53: 2.66, 54: 0, 55: 0.79, 56: 0.89, 57: 1.1, 58: 1.12, 59: 1.13, 60: 1.14, 61: 1.13, 62: 1.17, 63: 1.2, 64: 1.2, 65: 1.2, 66: 1.22, 67: 1.23, 68: 1.24, 69: 1.25, 70: 1.1, 71: 1.27, 72: 1.3, 73: 1.5, 74: 2.36, 75: 1.9, 76: 2.2, 77: 2.2, 78: 2.28, 79: 2.54, 80: 2, 81: 2.04, 82: 2.33, 83: 2.02, 84: 2, 85: 2.2, 86: 0, 87: 0.7, 88: 0.9, 89: 1.1, 90: 1.3, 91: 1.5, 92: 1.38, 93: 1.36, 94: 1.28, 95: 1.3, 96: 1.3, 97: 1.3, 98: 1.3, 99: 1.3, 100: 1.3, 101: 1.3, 102: 1.3, 103: 10.0, 104: 10.0, 105:10.0, 106: 10.0, 107: 10.0, 108: 10.0, 109: 10.0}
         Elec_aff={1: 0.75420375, 2: 0.0, 3: 0.618049, 4: 0.0, 5: 0.279723, 6: 1.262118, 7: -0.07, 8: 1.461112, 9: 3.4011887, 10: 0.0, 11: 0.547926, 12: 0.0, 13: 0.43283, 14: 1.389521, 15: 0.7465, 16: 2.0771029, 17: 3.612724, 18: 0.0, 19: 0.501459, 20: 0.02455, 21: 0.188, 22: 0.084, 23: 0.525, 24: 0.67584, 25: 0.0, 26: 0.151, 27: 0.6633, 28: 1.15716, 29: 1.23578, 30: 0.0, 31: 0.41, 32: 1.232712, 33: 0.814, 34: 2.02067, 35: 3.363588, 36: 0.0, 37: 0.485916, 38: 0.05206, 39: 0.307, 40: 0.426, 41: 0.893, 42: 0.7472, 43: 0.55, 44: 1.04638, 45: 1.14289, 46: 0.56214, 47: 1.30447, 48: 0.0, 49: 0.404, 50: 1.112066, 51: 1.047401, 52: 1.970875, 53: 3.059038, 54: 0.0, 55: 0.471626, 56: 0.14462, 57: 0.47, 58: 0.5, 59: 0.5, 60: 0.5, 61: 0.5, 62: 0.5, 63: 0.5, 64: 0.5, 65: 0.5, 66: 0.5, 67: 0.5, 68: 0.5, 69: 0.5, 70: 0.5, 71: 0.5, 72: 0.0, 73: 0.322, 74: 0.815, 75: 0.15, 76: 1.0778, 77: 1.56436, 78: 2.1251, 79: 2.30861, 80: 0.0, 81: 0.377, 82: 0.364, 83: 0.942363, 84: 1.9, 85: 2.8, 86: 0.0, 87: 0.0, 88: 0.0, 89: 0.0, 90: 0.0, 91: 0.0, 92: 0.0, 93: 0.0, 94: 0.0, 95: 0.0, 96: 0.0, 97: 0.0, 98: 0.0, 99: 0.0, 100: 0.0, 101: 0.0, 102: 0.0, 103: 0.0, 104: 0.0, 105: 0.0, 106: 0.0, 107: 0.0, 108: 0.0, 109: 0.0}
         Ion_nrg={1: 13.5984, 2: 24.5874, 3: 5.3917, 4: 9.3227, 5: 8.298, 6: 11.2603, 7: 14.5341, 8: 13.6181, 9: 17.4228, 10: 21.5645, 11: 5.1391, 12: 7.6462, 13: 5.9858, 14: 8.1517, 15: 10.4867, 16: 10.36, 17: 12.9676, 18: 15.7596, 19: 4.3407, 20: 6.1132, 21: 6.5615, 22: 6.8281, 23: 6.7462, 24: 6.7665, 25: 7.434, 26: 7.9024, 27: 7.881, 28: 7.6398, 29: 7.7264, 30: 9.3942, 31: 5.9993, 32: 7.8994, 33: 9.7886, 34: 9.7524, 35: 11.8138, 36: 13.9996, 37: 4.1771, 38: 5.6949, 39: 6.2173, 40: 6.6339, 41: 6.7589, 42: 7.0924, 43: 7.28, 44: 7.3605, 45: 7.4589, 46: 8.3369, 47: 7.5762, 48: 8.9938, 49: 5.7864, 50: 7.3439, 51: 8.6084, 52: 9.0096, 53: 10.4513, 54: 12.1298, 55: 3.8939, 56: 5.2117, 57: 5.5769, 58: 5.5387, 59: 5.473, 60: 5.525, 61: 5.582, 62: 5.6437, 63: 5.6704, 64: 6.1498, 65: 5.8638, 66: 5.9389, 67: 6.0215, 68: 6.1077, 69: 6.1843, 70: 6.2542, 71: 5.4259, 72: 6.8251, 73: 7.5496, 74: 7.864, 75: 7.8335, 76: 8.4382, 77: 8.967, 78: 8.9588, 79: 9.2255, 80: 10.4375, 81: 6.1082, 82: 7.4167, 83: 7.2855, 84: 8.414, 85: -1, 86: 10.7485, 87: 4.0727, 88: 5.2784, 89: 5.17, 90: 6.3067, 91: 5.89, 92: 6.1941, 93: 6.2657, 94: 6.026, 95: 5.9738, 96: 5.9914, 97: 6.1979, 98: 6.2817, 99: 6.42, 100: 6.5, 101: 6.58, 102: 6.65, 103: 4.9, 104: 6.0, 105: -1, 106: -1, 107: -1, 108: -1, 109: -1}
         DEN=Elec_neg[sa] -Elec_neg[sb]
         DEA=Elec_aff[sa] -Elec_aff[sb]
         DIE=Ion_nrg[sa] - Ion_nrg[sb]
         p=1
         if self.deltaEN>0: p*=np.exp(-0.5*(DEN/self.deltaEN)**2)
         if self.deltaEA>0: p*=np.exp(-0.5*(DEA/self.deltaEA)**2)
         if self.deltaIE>0: p*=np.exp(-0.5*(DIE/self.deltaIE)**2)
         return p
   def __init__(self, rules={},deltaEN=1e100, deltaEA=1e100, deltaIE=1e100, mu=0):
      self.rules = rules.copy()
      self.mu = mu
      self.deltaEA= deltaEA
      self.deltaIE= deltaIE
      self.deltaEN= deltaEN

def main(deltaEN,deltaEA,deltaIE,splist,alchemyrules):
   if (alchemyrules=="none"):
      alchem=alchemy_mendeleev(deltaEN=deltaEN,deltaIE=deltaIE,deltaEA=deltaEA,rules={})
   else:
       r=alchemyrules.replace('"', '').strip()
       r=alchemyrules.replace("'", '').strip()
       r=ast.literal_eval(r)
       print >> sys.stderr, "Using Alchemy rules: ", r,"\n"
       alchem=alchemy_mendelev(deltaEN=deltaEN,deltaIE=deltaIE,deltaEA=deltaEA,rules=r)
   rule={}
   for sa in splist:
      for sb in splist:
        if (sb >sa):
           rule[(sa,sb)]=alchem.getpair(sa,sb)
   print rule
   f="alchemy"
   if deltaEN > 0: f+="_dEN"+str(deltaEN)
   if deltaEA > 0: f+="_dEA"+str(deltaEA)
   if deltaIE > 0: f+="_dIE"+str(deltaIE)
   f+=".pickle"
   file = open(f,"wb")
   gc.disable()
   pickle.dump(rule, file,protocol=pickle.HIGHEST_PROTOCOL) # HIGHEST_PROTOCOL is 2 in py 2.7
   file.close()
   gc.enable()
if __name__ == '__main__':
      parser = argparse.ArgumentParser(description="""Utility code to generate alchemy rules based on atomic properties.""")

      parser.add_argument("--deltaEN",type=float,default="-1", help="Delta value for electronegativity")
      parser.add_argument("--deltaEA",type=float,default="-1", help="Delta value for electron affinity [eV]")
      parser.add_argument("--deltaIE",type=float,default="-1", help="Delta value for ionization energy [eV]")
      parser.add_argument("--species", type=str, default="all", help="list of species (e.g. --species 1,6  for H and C )")
      parser.add_argument("--rules", type=str, default="none", help='Dictionary-style rule specification in quote (e.g. --rules "{(6,7):1,(6,8):1}"')
      args = parser.parse_args()
      if args.species == "all":
         splist = []
         for i in range(1,109):
            splist.append(i)
      else:
         splist = sorted(map(int,args.species.split(',')))
      main(deltaEN=args.deltaEN,deltaIE=args.deltaIE,deltaEA=args.deltaEA,splist=splist,alchemyrules=args.rules)


#!/usr/bin/env python
'''
This script performs the t-test validation for non-bit-for-bit results for the
CICE model.

Written by: Matthew Turner
Date: October, 2017
'''

import os
import sys
import logging
import numpy as np
import numpy.ma as ma
import netCDF4 as nc

def maenumerate(marr):
    '''
    This function provides the enumerate functionality for masked arrays
    '''
    mask = ~marr.mask.ravel()
    try:   # Python 2
        import itertools
        for i, m in itertools.izip(np.ndindex(marr.shape[-2:]), mask):
            if m: yield i
    except:  # Python 3
        for i, m in zip(np.ndindex(marr.shape[-2:]), mask):
            if m: yield i

def gen_filenames(base_dir, test_dir):
    '''
    This function is passed the directories of the baseline and test history
    files and generates a list of filenames for each.
    '''
    # The path to output files for simulation 'a' (the '-bc' simulation)
    if base_dir.endswith(('history', 'history/')):
        path_a = base_dir
    else:
        path_a = base_dir + '/history/'

    # The path to output files for simulation 'b' (the test simulation)
    if test_dir.endswith(('history', 'history/')):
        path_b = test_dir
    else:
        path_b = test_dir + '/history/'

    # Find the number of output files to be read in
    files_a = [i for i in os.listdir(path_a+'/') if i.startswith('iceh_inst.')]
    files_b = [i for i in os.listdir(path_b+'/') if i.startswith('iceh_inst.')]

    if not len(files_a) == len(files_b):
        logger.error("Number of output files for baseline simulation does not match the number" + \
              " of files for the test simulation.  Exiting...\n" + \
              "Baseline directory: {}\n".format(path_a) + \
              "   # of files: {}\n".format(len(files_a)) + \
              "Test directory: {}\n".format(path_b) + \
              "   # of files: {}".format(len(files_b)))
        sys.exit(-1)

    logger.info("Number of files: %d", len(files_a))

    return path_a, path_b, files_a, files_b

def get_geom(path, file):
    '''
    This function reads the ni, nj, tlat, and tlon variables from a netcdf file
    '''
    fid = nc.Dataset("{}/{}".format(path, file), 'r')
    tlat = fid.variables['TLAT'][:]
    tlon = fid.variables['TLON'][:]
    ni = fid.dimensions['ni'].size
    nj = fid.dimensions['nj'].size
    fid.close()

    return ni, nj, tlat, tlon

def read_data(path_a, path_b, files_a, files_b, ni, nj):
    '''
    Read the baseline and test data for sea ice thickness.  The calculate
    the difference for all locations where sea ice thickness is greater
    than 0.01 meters.
    '''
    def fill_data_array(path, files, nj, ni):
        '''Function to fill the data arrays'''
        # Initialize the data array
        data = np.zeros((len(files), nj, ni),dtype=np.float32)
        # Read in the data
        logger.debug('Reading in data for files in %s', path)
        cnt = 0
        for fname in sorted(files):
            nfid = nc.Dataset("{}/{}".format(path, fname), 'r')
            fill_value = nfid.variables[var]._FillValue
            data[cnt, :, :] = nfid.variables[var][:]
            cnt += 1
            nfid.close()
        data[data == fill_value] = 0.0

        return data

    def calc_diff(data_a, data_b):
        '''
        Calculate the difference and mask the points where the the difference at
        every timestep is 0, or the sea ice thickness for every timestep is < 0.01 meters
        for data_a or data_b
        '''
        data_d = data_a - data_b
        mask_d = np.logical_or(\
                      np.logical_or(\
                           np.all(np.equal(data_d, 0.), axis=0), np.all(data_a < 0.01, axis=0))\
                      , np.all(data_b < 0.01, axis=0))
        mask_array_a = np.zeros_like(data_d)

        for x, value in np.ndenumerate(mask_d):
            i, j = x
            mask_array_a[:, i, j] = value
        del mask_d

        data_a = ma.masked_array(data_a, mask=mask_array_a)
        data_b = ma.masked_array(data_b, mask=mask_array_a)
        data_d = ma.masked_array(data_d, mask=mask_array_a)

        del mask_array_a

        return data_a, data_b, data_d

    var = 'hi'

    data_a = fill_data_array(path_a, files_a, nj, ni)
    data_b = fill_data_array(path_b, files_b, nj, ni)

    data_a, data_b, data_d = calc_diff(data_a, data_b)

    return data_a, data_b, data_d

def two_stage_test(data_a, num_files, data_d, fname, path):
    '''
    This function performs the Two-Stage Paired Thickness Test
    '''
    def stage_one(data_d, num_files, mean_d, variance_d):
        logger.debug('Running step 1 of 2-stage test')

        # Calculate the mean from 1:end-1 and 2:end
        mean_nm1_d = np.mean(data_d[:-1, :, :], axis=0)
        mean_2n_d = np.mean(data_d[1:, :, :], axis=0)

        # Calculate equation (5) for both simulations
        r1_num = np.zeros_like(mean_d)
        r1_den1 = np.zeros_like(mean_d)
        r1_den2 = np.zeros_like(mean_d)
        for i in np.arange(np.size(data_a, axis=0)-1):
            r1_num = r1_num + (data_d[i, :, :]-mean_nm1_d[:, :])*(data_d[i+1, :, :]-mean_2n_d[:, :])
            r1_den1 = r1_den1 + np.square(data_d[i, :, :]-mean_nm1_d[:, :])

        for i in np.arange(1, np.size(data_a, axis=0)):
            r1_den2 = r1_den2 + np.square(data_d[i, :, :] - mean_2n_d[:, :])

        r1 = r1_num / np.sqrt(r1_den1*r1_den2)

        # Calculate the effective sample size
        n_eff = num_files * ((1.-r1) / (1.+r1))
        n_eff[n_eff < 2] = 2
        n_eff[n_eff > num_files] = num_files

        # Calculate the t-statistic with n_eff
        t_val = mean_d / np.sqrt(variance_d / n_eff)

        # Effective degrees of freedom
        df = n_eff - 1

        # Read in t_crit table
        nfid = nc.Dataset("configuration/scripts/tests/QC/CICE_t_critical_p0.8.nc", 'r')
        df_table = nfid.variables['df'][:]
        t_crit_table = nfid.variables['tcrit'][:]
        nfid.close()
        t_crit = np.zeros_like(t_val)

        # Calculate critical t-value for each grid cell, based on the t_crit table
        for x in maenumerate(data_d):
            min_val = np.min(np.abs(df[x]-df_table))
            idx = np.where(np.abs(df[x]-df_table) == min_val)
            # Handle the cases where the data point falls exactly half way between
            # 2 critical T-values (i.e., idx has more than 1 value in it)
            while True:
                try:
                    idx = idx[0]
                except:
                    break
            t_crit[x] = t_crit_table[idx]

        # Create an array of Pass / Fail values for each grid cell
        H1 = np.abs(t_val) > t_crit

        return n_eff, H1, r1, t_crit

    # Calculate the mean of the difference
    mean_d = np.mean(data_d, axis=0)

    # Loop through each timestep and calculate the square of the difference.
    # This is required (instead of just np.square(data_d - mean_d) to reduce
    # the memory footprint of the script.
    tmp1 = np.zeros_like(data_d)
    for i in np.arange(np.shape(data_d)[0]):
        tmp1[i,:,:] = np.square(data_d[i,:,:] - mean_d[:,:])
    variance_d = np.sum(tmp1) / float(num_files - 1)

    n_eff, H1, r1, t_crit = stage_one(data_d, num_files, mean_d, variance_d)

    if np.all(H1 == False) and np.all(n_eff >= 30):
        # H0 confirmed in all cells, and all effective sample size >= 30
        logger.debug('H0 confirmed in all cells, and all effective sample size >= 30')
        logger.info('2 stage test passed')
        return True, H1

    elif np.all(H1):
        # H1 in every grid cell
        logger.debug('H1 in all cells')
        logger.info('2 stage test failed')
        return False, H1

########### H0 confirmed for some grid cells with n_eff < 30 ############
    logger.debug('Number of H1 grid cells after stage 1 = %d', np.sum(H1))

    logger.debug('Running step 2 of 2-stage test')

    # Find the indices where n_eff is less than 30, and H0 is confirmed
    tmp_idx = np.where(n_eff < 30) and np.where(H1 == False)

    # Calculate the T-statistic using actual sample size
    t_val = mean_d / np.sqrt(variance_d / num_files)

    # Find t_crit from the nearest value on the Lookup Table Test
    nfid = nc.Dataset("configuration/scripts/tests/QC/CICE_Lookup_Table_p0.8_n1825.nc", 'r')
    r1_table = nfid.variables['r1'][:]
    t_crit_table = nfid.variables['tcrit'][:]
    nfid.close()

    # Fill t_crit based on lookup table
    for x in maenumerate(data_d):
        min_val = np.min(np.abs(r1[x]-r1_table))
        idx = np.where(np.abs(r1[x]-r1_table) == min_val)
        # Handle the cases where the data point falls exactly half way between
        # 2 critical T-values (i.e., idx has more than 1 value in it)
        while True:
            try:
                idx = idx[0]
            except:
                break
        t_crit[x] = t_crit_table[idx]

    # Create an array showing locations of Pass / Fail grid cells
    H1[tmp_idx] = abs(t_val[tmp_idx]) > t_crit[tmp_idx]

    logger.debug('Number of H1 grid cells after stage 2 = %d', np.sum(H1))

    if np.all(H1):
        # H1 in all grid cells
        logger.debug('H1 in all cells, stage 2')
        logger.info('2 Stage Test Failed')
        return False, H1

    elif np.all(H1 == False):
        # H0 confirmed in all grid cells
        logger.debug('H0 confirmed in all cells with n_eff < 30')
        logger.info('2 Stage Test Passed')
        return True, H1

####### Some grid cells have H0 confirmed, and some do not ######

    # Calculate the area-weighted fraction of the test region that failed (f_val).
    #   If f_val is greater than or equal to the critical fraction, the test fails"
    f_val = critical_fraction(data_a, H1, fname, path)
    f_crit = 0.5
    if f_val >= f_crit:
        logger.info('2 Stage Test Failed')
        logger.debug('Area-weighted fraction of failures is greater than ' + \
                    'critical fraction.  Test failed.')
        logger.debug('Area-weighted fraction of failures = %f', f_val)
        return False, H1
    else:
        logger.info('2 Stage Test Passed')
        logger.debug('Area-weighted fraction of failures = %f', f_val)
        return True, H1

def critical_fraction(data_a, failures, fname, path_a):
    '''
    This function calculates the area-weighted average of cells where H1 is true.
    '''
    logger.debug('Calculating area-weighted average of H1 cells')
    # First calculate the weight attributed to each grid point (based on Area)
    nfid = nc.Dataset("{}/{}".format(path_a, fname), 'r')
    tarea = nfid.variables['tarea'][:]
    nfid.close()
    tarea = ma.masked_array(tarea, mask=data_a[0, :, :].mask)
    area_weight = tarea / np.sum(tarea)

    # Calculate the area weight of the failing grid cells
    weight_tot = 0
    weight_fail = 0
    for x in maenumerate(data_a):
        weight_tot += area_weight[x]
        if failures[x]:
            weight_fail += area_weight[x]

    return weight_fail/weight_tot

def skill_test(path_a, fname, data_a, data_b, num_files, hemisphere):
    '''Calculate Taylor Skill Score'''
    # First calculate the weight attributed to each grid point (based on Area)
    nfid = nc.Dataset("{}/{}".format(path_a, fname), 'r')
    tarea = nfid.variables['tarea'][:]
    nfid.close()
    tarea = ma.masked_array(tarea, mask=data_a[0, :, :].mask)
    area_weight = tarea / np.sum(tarea)

    weighted_mean_a = 0
    weighted_mean_b = 0
    for i in np.arange(num_files):
        weighted_mean_a = weighted_mean_a + np.sum(area_weight*data_a[i, :, :])
        weighted_mean_b = weighted_mean_b + np.sum(area_weight*data_b[i, :, :])

    weighted_mean_a = weighted_mean_a / num_files
    weighted_mean_b = weighted_mean_b / num_files

    nonzero_weights = np.count_nonzero(area_weight)
    area_var_a = 0
    area_var_b = 0
    for t in np.arange(num_files):
        area_var_a = area_var_a + np.sum(area_weight*np.square(data_a[t, :, :]-weighted_mean_a))
        area_var_b = area_var_b + np.sum(area_weight*np.square(data_b[t, :, :]-weighted_mean_b))

    area_var_a = nonzero_weights / (num_files * nonzero_weights - 1.) * area_var_a
    area_var_b = nonzero_weights / (num_files * nonzero_weights - 1.) * area_var_b
    std_a = np.sqrt(area_var_a)
    std_b = np.sqrt(area_var_b)

    combined_cov = 0
    for i in np.arange(num_files):
        combined_cov = combined_cov + np.sum(area_weight*(data_a[i, :, :]-weighted_mean_a)*\
                                            (data_b[i, :, :]-weighted_mean_b))

    combined_cov = nonzero_weights / (num_files * nonzero_weights - 1.) * combined_cov

    weighted_r = combined_cov / (std_a*std_b)

    s = np.square((1+weighted_r)*(std_a*std_b)/\
                 (area_var_a + area_var_b))

    logger.debug('%s Hemisphere skill score = %f', hemisphere, s)

    s_crit = 0.99
    if s < 0 or s > 1:
        logger.error('Skill score out of range for %s Hemisphere', hemisphere)
        return False
    elif s > s_crit:
        logger.info('Quadratic Skill Test Passed for %s Hemisphere', hemisphere)
        return True
    else:
        logger.info('Quadratic Skill Test Failed for %s Hemisphere', hemisphere)
        return False

def plot_data(data, lat, lon, units, case, plot_type):
    '''This function plots CICE data and creates a .png file.'''

    try:
        # Load the necessary plotting libraries
        import matplotlib.pyplot as plt
        from mpl_toolkits.basemap import Basemap
        from mpl_toolkits.axes_grid1 import make_axes_locatable
    except ImportError:
        logger.warning('Error loading necessary Python modules in plot_data function')
        return

    # Suppress Matplotlib deprecation warnings
    import warnings
    warnings.filterwarnings("ignore", category=UserWarning)

    # Create the figure and axis
    fig, axes = plt.subplots(nrows=1, ncols=2,figsize=(14, 8))

    # Plot the northern hemisphere data as a scatter plot
    # Create the basemap, and draw boundaries
    plt.sca(axes[0])
    m = Basemap(projection='npstere', boundinglat=35,lon_0=270, resolution='l')
    m.drawcoastlines()
    m.fillcontinents()
    m.drawcountries()

    if plot_type == 'scatter':
        x, y = m(lon,lat)
        sc = m.scatter(x, y, c=data, cmap='jet', lw=0, s=4)
    else:
        # Create new arrays to add 1 additional longitude value to prevent a
        # small amount of whitespace around longitude of 0/360 degrees.
        lon_cyc = np.zeros((lon.shape[0],lon.shape[1]+1))
        mask = np.zeros((data.shape[0],data.shape[1]+1))
        lat_cyc = np.zeros((lat.shape[0],lat.shape[1]+1))

        mask[:,0:-1] = data.mask[:,:]
        mask[:,-1] = data.mask[:,0]
        lon_cyc[:,0:-1] = lon[:,:]; lon_cyc[:,-1] = lon[:,0]
        lat_cyc[:,0:-1] = lat[:,:]; lat_cyc[:,-1] = lat[:,0]

        lon1 = np.ma.masked_array(lon_cyc, mask=mask)
        lat1 = np.ma.masked_array(lat_cyc, mask=mask)

        d = np.zeros((data.shape[0],data.shape[1]+1))
        d[:,0:-1] = data[:,:]
        d[:,-1] = data[:,0]
        d1 = np.ma.masked_array(d,mask=mask)

        x, y = m(lon1.data, lat1.data)

        if plot_type == 'contour':
            sc = m.contourf(x, y, d1, cmap='jet')
        else:  # pcolor
            sc = m.pcolor(x, y, d1, cmap='jet')

    m.drawparallels(np.arange(-90.,120.,15.),labels=[1,0,0,0]) # draw parallels
    m.drawmeridians(np.arange(0.,420.,30.),labels=[1,1,1,1]) # draw meridians

    # Plot the southern hemisphere data as a scatter plot
    plt.sca(axes[1])
    m = Basemap(projection='spstere', boundinglat=-45,lon_0=270, resolution='l')
    m.drawcoastlines()
    m.fillcontinents()
    m.drawcountries()

    if plot_type == 'scatter':
        x, y = m(lon,lat)
        sc = m.scatter(x, y, c=data, cmap='jet', lw=0, s=4)
    else:
        x, y = m(lon1.data, lat1.data)

        # Bandaid for a bug in the version of Basemap used during development
        outside = (x <= m.xmin) | (x >= m.xmax) | (y <= m.ymin) | (y >= m.ymax)
        tmp = np.ma.masked_where(outside,d1)

        if plot_type == 'contour':
            sc = m.contourf(x, y, tmp, cmap='jet')
        else:  # pcolor
            sc = m.pcolor(x, y, tmp, cmap='jet')

    m.drawparallels(np.arange(-90.,120.,15.),labels=[1,0,0,0]) # draw parallels
    m.drawmeridians(np.arange(0.,420.,30.),labels=[1,1,1,1]) # draw meridians

    plt.suptitle('CICE Mean Ice Thickness\n{}'.format(case), y=0.95)

    # Make some room at the bottom of the figure, and create a colorbar
    fig.subplots_adjust(bottom=0.2)
    cbar_ax = fig.add_axes([0.11,0.1,0.8,0.05])
    if '\n- ' in case:
      # If making a difference plot, use scientific notation for colorbar
      cb = plt.colorbar(sc, cax=cbar_ax, orientation="horizontal", format="%.2e")
    else:
      # If plotting non-difference data, do not use scientific notation for colorbar
      cb = plt.colorbar(sc, cax=cbar_ax, orientation="horizontal", format="%.2f")
    cb.set_label(units, x=1.0)

    outfile = 'ice_thickness_{}.png'.format(case.replace('\n- ','_minus_'))
    logger.info('Creating map of the data ({})'.format(outfile))
    plt.savefig(outfile, dpi=300, bbox_inches='tight')

def plot_two_stage_failures(data, lat, lon):
    '''This function plots each grid cell and whether or not it Passed or Failed
       the two-stage test.  It then either creates a .png file
       (two_stage_test_failure_map.png), or saves the failure locations to a
       text file.
    '''

    # Convert the boolean array (data) to an integer array
    int_data = data.astype(int)

    try:
        logger.info('Creating map of the failures (two_stage_test_failure_map.png)')
        # Load the necessary plotting libraries
        import matplotlib.pyplot as plt
        from mpl_toolkits.basemap import Basemap
        from mpl_toolkits.axes_grid1 import make_axes_locatable
        from matplotlib.colors import LinearSegmentedColormap

        # Suppress Matplotlib deprecation warnings
        import warnings
        warnings.filterwarnings("ignore", category=UserWarning)

        # Create the figure and axis
        fig = plt.figure(figsize=(12, 8))
        ax = fig.add_axes([0.05, 0.08, 0.9, 0.9])

        # Create the basemap, and draw boundaries
        m = Basemap(projection='moll', lon_0=0., resolution='l')
        m.drawmapboundary(fill_color='white')
        m.drawcoastlines()
        m.drawcountries()

        # Create the custom colormap
        colors = [(0, 0, 1), (1, 0, 0)]  # Blue, Red
        cmap_name = 'RB_2bins'
        cm = LinearSegmentedColormap.from_list(cmap_name, colors, N=2)

        # Plot the data as a scatter plot
        x, y = m(lon, lat)
        sc = m.scatter(x, y, c=int_data, cmap=cm, lw=0, vmin=0, vmax=1, s=4)

        m.drawmeridians(np.arange(0, 360, 60), labels=[0, 0, 0, 1], fontsize=10)
        m.drawparallels(np.arange(-90, 90, 30), labels=[1, 0, 0, 0], fontsize=10)

        plt.title('CICE Two-Stage Test Failures')

        # Create the colorbar and add Pass / Fail labels
        divider = make_axes_locatable(ax)
        cax = divider.append_axes("bottom", size="5%", pad=0.5)
        cb = plt.colorbar(sc, cax=cax, orientation="horizontal", format="%.0f")
        cb.set_ticks([])
        cb.ax.text(-0.01, -0.5, 'PASS')
        cb.ax.text(0.99, -0.5, 'FAIL')

        plt.savefig('two_stage_test_failure_map.png', dpi=300)
    except:
        logger.warning('')
        logger.warning('Unable to plot the data.  Saving latitude and longitude')
        logger.warning('for ONLY failures to two_stage_failure_locations.txt')

        # Create a file and write the failures only to the file
        f = open('two_stage_failure_locations.txt', 'w')
        f.write('# CICE Two-stage test failures\n')
        f.write('# Longitude,Latitude\n')
        for i in range(data.shape[0]):
            for j in range(data.shape[1]):
                if (not data.mask[i, j]) and data[i, j]:
                    f.write('{},{}\n'.format(lon[i, j], lat[i, j]))

        f.close()

def main():
    import argparse
    parser = argparse.ArgumentParser(description='This script performs the T-test for \
                           CICE simulations that should be bit-for-bit, but are not.')
    parser.add_argument('base_dir', \
                help='Path to the baseline history (iceh_inst*) files.  REQUIRED')
    parser.add_argument('test_dir', \
                help='Path to the test history (iceh_inst*) files.  REQUIRED')
    parser.add_argument('-v', '--verbose', dest='verbose', help='Print debug output?', \
                        action='store_true')
    parser.add_argument('-pt','--plot_type', dest='plot_type', help='Specify type of plot \
                        to create', choices=['scatter','contour','pcolor'])

    parser.set_defaults(verbose=False)
    parser.set_defaults(plot_type='pcolor')

    # If no arguments are provided, print the help message
    if len(sys.argv) == 1:
        parser.print_help()
        sys.exit(1)

    args = parser.parse_args()

    # Set up the logger
    global logger
    if args.verbose:
        logging.basicConfig(level=logging.DEBUG)
    else:
        logging.basicConfig(level=logging.INFO)
    # Log to log file as well as stdout
    fh = logging.FileHandler(r'qc_log.txt', 'w')
    logger = logging.getLogger(__name__)
    logger.addHandler(fh)

    logger.info('Running QC test on the following directories:')
    logger.info('  {}'.format(args.base_dir))
    logger.info('  {}'.format(args.test_dir))

    dir_a, dir_b, files_base, files_test = gen_filenames(args.base_dir, args.test_dir)

    nfiles = len(files_base)

    nlon, nlat, t_lat, t_lon = get_geom(dir_a, files_base[0])

    data_base, data_test, data_diff = read_data(dir_a, dir_b, files_base, files_test, nlon, nlat)

    if np.ma.all(data_diff.mask):
        logger.info("Data is bit-for-bit.  No need to run QC test")
        sys.exit(0)

    # Run the two-stage test
    PASSED, H1_array = two_stage_test(data_base, nfiles, data_diff, files_base[0], dir_a)

    # Delete arrays that are no longer necessary
    del data_diff

    # If test failed, attempt to create a plot of the failure locations
    if not PASSED:
        plot_two_stage_failures(H1_array, t_lat, t_lon)

        # Create plots of mean ice thickness
        baseDir = os.path.abspath(args.base_dir).rstrip('history/').rstrip(\
                                                        'history').split('/')[-1]
        testDir = os.path.abspath(args.test_dir).rstrip('history/').rstrip( \
                                                        'history').split('/')[-1]
        plot_data(np.mean(data_base,axis=0), t_lat, t_lon, 'm', baseDir, args.plot_type)
        plot_data(np.mean(data_test,axis=0), t_lat, t_lon, 'm', testDir, args.plot_type)
        plot_data(np.mean(data_base-data_test,axis=0), t_lat, t_lon, 'm', '{}\n- {}'.\
                  format(baseDir,testDir), args.plot_type)

        logger.error('Quality Control Test FAILED')
        sys.exit(-1)

    # Create a northern hemisphere and southern hemisphere mask
    mask_tlat = t_lat < 0
    mask_nh = np.zeros_like(data_base)
    mask_sh = np.zeros_like(data_base)
    for (a, b), val in np.ndenumerate(mask_tlat):
        mask_nh[:, a, b] = val
        mask_sh[:, a, b] = not val

    # Run skill test on northern hemisphere
    data_nh_a = ma.masked_array(data_base, mask=mask_nh)
    data_nh_b = ma.masked_array(data_test, mask=mask_nh)
    if np.ma.all(data_nh_a.mask) and np.ma.all(data_nh_b.mask):
        logger.info("Northern Hemisphere data is bit-for-bit")
        PASSED_NH = True
    else:
        PASSED_NH = skill_test(dir_a, files_base[0], data_nh_a, data_nh_b, nfiles, 'Northern')

    # Run skill test on southern hemisphere
    data_sh_a = ma.masked_array(data_base, mask=mask_sh)
    data_sh_b = ma.masked_array(data_test, mask=mask_sh)
    if np.ma.all(data_sh_a.mask) and np.ma.all(data_sh_b.mask):
        logger.info("Southern Hemisphere data is bit-for-bit")
        PASSED_SH = True
    else:
        PASSED_SH = skill_test(dir_a, files_base[0], data_sh_a, data_sh_b, nfiles, 'Southern')

    PASSED_SKILL = PASSED_NH and PASSED_SH

    # Plot the ice thickness data for the base and test cases
    baseDir = os.path.abspath(args.base_dir).rstrip('history/').rstrip( \
                                                    'history').split('/')[-1]
    testDir = os.path.abspath(args.test_dir).rstrip('history/').rstrip( \
                                                    'history').split('/')[-1]
    plot_data(np.mean(data_base,axis=0), t_lat, t_lon, 'm', baseDir, args.plot_type)
    plot_data(np.mean(data_test,axis=0), t_lat, t_lon, 'm', testDir, args.plot_type)
    plot_data(np.mean(data_base-data_test,axis=0), t_lat, t_lon, 'm', '{}\n- {}'.\
              format(baseDir,testDir), args.plot_type)

    logger.info('')
    if not PASSED_SKILL:
        logger.error('Quality Control Test FAILED')
        sys.exit(1)  # exit with an error return code
    else:
        logger.info('Quality Control Test PASSED')
        sys.exit(0)  # exit with successfull return code

if __name__ == "__main__":
    main()