import sys
import os
import os.path
import copy
import collections
import ast

import numpy
import numpy as np
import pandas as pd

from limatix.lm_units import parseunits,multiplyunits,printunits

from matplotlib import pyplot as pl
#import pyximport
#pyximport.install()
from . import convolution
from .convolution import impulse_response,convolution_evaluation

#
#
# Inputs:
# Welder model: Vector A of complex amplitude
#               Vector alpha of complex exponents
#               time delay (syn_time_delay)
#               initial_displacement (should be -rho*c_A) because motion positive toward vibrometer
# Welder model predicts response to an impulse force away from the vibrometer
# with amplitude 1 N*s

# Also need a model for the pneumatic cylinder behaviour
# That provides the restoring force
# and the long-time behavior where the 1 N-s impulse
# accelerate the entire welder.

#
#
# Welder & Specimen model
# Represent as impulse response, possibly followed by continued ringing
# of one or more decaying sinusoids (i.e. exp(-alpha*t)

class dummy(object):
    pass

pkgpath = sys.modules[dummy.__module__].__file__
pkgdir=os.path.split(pkgpath)[0]

# All welder responses now loaded from welder_modeling_out.py
welder_model_output_dir = pkgdir

# Welder model provides globals welder_tip_tip_resp, welder_tip_elec_resp, welder_elec_tip_resp,
# and welder_elec_elec_resp
exec(open(os.path.join(welder_model_output_dir,"welder_modeling_out.py")).read(),globals())

## Temporarily zero out welder_tip_tip_resp for convergence testing
#welder_tip_tip_resp.h[:]=0.0
#welder_tip_tip_resp.h[:10]=-1e-28
#welder_tip_tip_resp.h[1]=-1e-28
#welder_tip_tip_resp.A[:]=0.0

default_welder_spring_constant=5000 # N/m -- bounciness of seals in welder pneumatic cylinder
default_R_contact=25.4e-3 # Hertzian contact parameter: 1 inch radius
default_welder_elec_freq=19890.0 # Frequency, Hz
default_dt=1e-6 # Time step, seconds

default_gpu_precision="single" # gpu_precision must be "single" or "double"

contact_F_nominal = 100.0 # N -- nominal contact force used to decide to neglect imaginary parts order of magnitude smaller
G_nominal = contact_F_nominal**(1.0/3.0)

#specimen_model_output_dir = '.'
#specimen_model_fname = "specimen_modeling_out.py"
#specimen_model_fname = "cantilever_modeling_out.py"
#specimen_model_fname = "cantilever_modeling_out_moredamping.py"
#specimen_model_fname = "cantilever_modeling_out_2019_10_01.py"
#specimen_model_fname = "cantilever_modeling_out_2019-11-09-finerseg.py.xz"

## response of specimen, evaluated at welder contact point
#specimen_dataframe = pd.read_csv(specimen_model_fname)
#specimen_resp =

#specimen_model_fpath=os.path.join(specimen_model_output_dir,specimen_model_fname)

def select_gpu_device(priority_list_str):
    """Based on a priority list string of the form:
     [
       ('NVIDIA CUDA','Quadro GP100'),
       ('Intel(R) OpenCL HD Graphics','Intel(R) Gen9 HD Graphics NEO'),
       ('Portable Computing Language','pthread-AMD EPYC 7351P 16-Core Processor')
     ]
    select the first of these devices found and return (context,device,queue).
    If priority_list_str=="" then None will be returned.
    Otherwise a ValueError or other error will be raised if the
    string is not parseable or if the device is not found.

    Each entry in the list is (Platform Name, Device Name)
    The available platforms and devices can be found by looking
    at the output of the "clinfo" command.
    """

    if priority_list_str != "":
        import pyopencl as cl
        gpu_device_priority_list = ast.literal_eval(priority_list_str)
        platforms=cl.get_platforms()
        platforms_byname = { platform.name: platform for platform in platforms }

        device = None

        for (gpu_platform_name,gpu_device_name) in gpu_device_priority_list:

            if gpu_platform_name in platforms_byname:
                platform = platforms_byname[gpu_platform_name]

                devices=platform.get_devices()
                devices_byname = { device.name: device for device in devices }

                if gpu_device_name in devices_byname:
                    device=devices_byname[gpu_device_name]
                    break
                pass
            pass

        if device is None:
            raise ValueError("No OpenCL devices found matching any entry in priority list %s. Use clinfo command to find platform name and device name. Priority list must be entered using list notation: \"[ ('platform1_name','device1_name'), ('platform2_name','device2_name') ]\"")
        context = cl.Context(devices=[device])
        queue = cl.CommandQueue(context)
        gpu_context_device_queue = (context,device,queue)
        pass
    else:
        gpu_context_device_queue=None
        pass
    return gpu_context_device_queue

def load_specimen_model(specimen_model_filepath):
    import pandas as pd
    specimen_dataframe = pd.read_csv(specimen_model_filepath,index_col="Time(s)")

    #dt=specimen_dataframe["Time(s)"][1]-specimen_dataframe["Time(s)"][0]
    dt=specimen_dataframe.index[1]-specimen_dataframe.index[0]

    assert(specimen_dataframe.index[0]==0.0)

    specimen_dict=collections.OrderedDict()
    specimen_units_dict=collections.OrderedDict()
    for column in specimen_dataframe.columns:
        if column=="Time(s)":
            continue
        unitstartidx = column.find("(")
        if unitstartidx < 0:
            raise ValueError("Dynamic model column name \"%s\" does not include units" % (column))
        fieldname=column[:unitstartidx]
        fieldunits_with_close_paren=column[(unitstartidx+1):].strip()
        if not fieldunits_with_close_paren.endswith(")"):
            raise ValueError("Dynamic model column \"%s\" units do not end with close parentheses" % (column))
        fieldunits = fieldunits_with_close_paren[:-1]

        specimen_dict[fieldname]=convolution.impulse_response(
            h=np.array(specimen_dataframe[column]),
            dt=dt,
            t0=0.0,
            A=np.array((0,),dtype='d'),
            D=0.0,
            alpha=np.array((1.0,),dtype='d'))

        specimen_units_dict[fieldname]=fieldunits
        pass

    return (specimen_dict,specimen_units_dict) # keys are "specimen_resp", etc.; values are impulse_response instances

def contact_model(specimen_dict,
                  specimen_units_dict,
                  t0_t1, # Excitation start time (seconds)
                  t2_t3, # Excitation end time (seconds)
                  t4, # Time to calculate out to (seconds)
                  mass_of_welder_and_slider, # mass, kg
                  pneumatic_force, # Force, N
                  welder_elec_ampl, # Amplitude (au...? should be volts)
                  specimen_E, specimen_nu, # Specimen elastic params
                  welder_spring_constant=default_welder_spring_constant, # N/m -- bounciness of seals in welder pneumatic cylinder
                  R_contact=default_R_contact, # Hertzian contact parameter: 1 inch radius
                  welder_elec_freq=default_welder_elec_freq, # Frequency, Hz
                  dt=default_dt, # Time step, seconds
                  gpu_context_device_queue=None,
                  gpu_precision=default_gpu_precision):

    # specimen_model defines specimen_resp, specimen_mobility, specimen_laser,
    # specimen_crackcenternormalstrain, and specimen_crackcentershearstrain

    ## temporarily hardwire gpu_precision to double
    #gpu_precision="double"
    #gpu_context_device_queue=None # Temporarily disable GPU

    #s(pecimen_dict,specimen_units_dict)=load_specimen_model(specimen_model_fpath)

    # IDEA: Add small uniform real part to specimen_resp in frequency domain
    # To keep the phase away from the edge.... FAILED

    # IDEA: Add damping to Hertzian contact spring

    # Resample specimen data to desired dt value
    for key in specimen_dict:

        specimen_dict[key].resample(dt)
        pass

    specimen_resp=specimen_dict["specimen_resp"]
    del specimen_dict["specimen_resp"] # Remove specimen_resp from specimen_dict so convolutions with entries in specimen_dict will not be redundant
    if "specimen_mobility" in specimen_dict:
        del specimen_dict["specimen_mobility"] # Don't care about predicting specimen contact velocity
        pass

    ## Temporarily zero out specimen_resp
    #specimen_resp.h[:]=0.0

    assert(welder_elec_tip_resp.h[0]==0.0) # response of tip to electrical excitation MUST be delayed
    assert(welder_elec_elec_resp.h[0]==0.0) # electrical response to electrical excitation MUST be delayed
    assert(welder_tip_elec_resp.h[0]==0.0) # electrical response to tip impulse MUST be delayed

    welder_tip_tip_resp_local=copy.deepcopy(welder_tip_tip_resp)
    welder_tip_elec_resp_local=copy.deepcopy(welder_tip_elec_resp)
    welder_elec_tip_resp_local=copy.deepcopy(welder_elec_tip_resp)
    welder_elec_elec_resp_local=copy.deepcopy(welder_elec_elec_resp)

    # resample all responses to desired timestep
    welder_tip_tip_resp_local.resample(dt)
    welder_tip_elec_resp_local.resample(dt)
    welder_tip_elec_resp_local.h[0]=0.0 # resampling can mess this up
    welder_elec_tip_resp_local.resample(dt)
    welder_elec_tip_resp_local.h[0]=0.0 # resampling can mess tihs up
    welder_elec_elec_resp_local.resample(dt)
    welder_elec_elec_resp_local.h[0]=0.0 # resampling can mess this up

    # welder_tip_tip_resp_local.h[0] is negative... specimen_resp.h[0] is positive,
    # (criteria for solving for contact force, below)
    assert(welder_tip_tip_resp_local.h[1] < 0.0)
    #assert(specimen_resp.h[0] > 0.0)

    # Connect specimen and welder models at contact
    #  * Sum of forces at contact = 0 (contact massless)
    #  * No overlap between specimen and welder.

    # z positive towards specimen
    welder_tip_z = 0
    specimen_z = 0

    #max_t = 0.1
    #max_t = 0.3
    #max_t=0.2
    trange = np.arange(t4/dt)*dt

    #mass_of_welder_and_slider=2.0   # mass, kg !!!*** Need to properly measure
    #pneumatic_force = 300 # Force, N  ***!!! Not necessarily representative
    # NOTE: If welder_spring_constant is nonzero, the pneumatic force ( in
    # equilbrium) be split between contact load and spring!!!

    # NOTE: with welder_overall_dashpot = 0
    # get resonance (~170Hz as of this writing), which
    # matches (1/(2pi))*sqrt(k/m) for
    # m = mass_of_welder_and_slider = 2 and
    # k = 1.0/(np.sum(specimen_resp.eval(np.arange(trange.shape[0])).real)*dt) = 2.5e6 N/m
    # represents the effective DC stiffness of the specimen
    #  (we should consider the welder and contact stiffnesses as well probably,
    #  but they are much stiffer!)
    # For critically damped, set damping ratio = c/(2sqrt(m*k)) = 1
    # so c = 2*sqrt(m*k) = 4472 for critically damped.
    # Probably want significantly underdamped, so set c=400 N/(m/s)
    welder_overall_dashpot = 1000 # dashpot coefficient simulating absorption of pneumatic cylinder

    # ... at welder 50% amplitude setting , expect open circuit velocity of roughly 80 um p-p at 20 kHz. this corresponds to a welder_elec_ampl of ~1.8e7 (UNITS?)
    #welder_elec_ampl= 1.8e7*1.0

    #... double it!
    #welder_elec_ampl= 1.8e7*5.0

    #welder_elec_freq=19890.0

    # Plot welder open-circuit behavior with:
    #
    #welder_elec_input = welder_elec_ampl*np.cos(2*np.pi*welder_elec_freq*trange)
    #welder_elec_tip_conv = convolution_evaluation.blank_from_imp_resp(welder_elec_tip_resp)
    #welder_tip_z = np.zeros(trange.shape,dtype='d')
    #for tcnt in range(trange.shape[0]):
    #    welder_tip_z[tcnt]=welder_elec_tip_conv.step(welder_elec_input[tcnt])
    #    pass
    #pl.plot(trange,welder_tip_z)

    # NOTE: This spring represents the limited contact zone,
    # and helps to prevent system resonances involving
    # high Q, high-mobility resonances like the 20 kHz welder motion
    #kspring = 120e9*(np.pi*12e-3**2/4.0)/(1e-3)*.1 # *100)  # contact stiffness, N/m
    # NOTE: kspring replaced by Hertzian contact model
    # Hertzian contact parameters
    #R_contact=25.4e-3 # 1 inch radius
    # Welder material (Ti)
    nu1=.342
    E1=113.8e9

    ## Al
    #nu2=.33
    #E2=68.9e9

    Estar = 1.0/( (1.0-nu1**2.0)/E1 + (1.0-specimen_nu**2.0)/specimen_E)

    # This spring represents the bounciness of the seals in the pneumatic cylinder... based on 20 ms resonant period, and assuming 2kg mass,
    # f=sqrt(k/m) = 1/.02
    # f*sqrt(m) = sqrt(k)
    # k = f^2*m = (1/.02)^2 * 2.0 = 5000 N/m
    # ... if you set this to 0, then you get instead pure
    # pneumatic cylinder behavior
    #welder_spring_constant = 5000

    welder_overall_velocity=0.0 # positive towards specimen
    #welder_overall_pos=-.5e-3 # half-mm distance initially # positive towards specimen
    #welder_overall_pos=0e-3 # no distance initially # positive towards specimen

    # This code is for starting from equilibrium displacement
    # Next line is correct for pneumatic cylinder but no spring
    #welder_overall_pos = convolution_evaluation.quiescent_value(pneumatic_force,specimen_resp.A,specimen_resp.alpha,specimen_resp.dt,specimen_resp.h).real - convolution_evaluation.quiescent_value(pneumatic_force,welder_tip_tip_resp_local.A,welder_tip_tip_resp_local.alpha,welder_tip_tip_resp_local.dt,welder_tip_tip_resp_local.h).real

    # overall_pos = specimen_quiescent_coeff*(pneumatic-welder_spring_constant*overall_pos) - welder_quiescent_coeff*(pneumatic-welder_spring_constant*overall_pos)
    # overall_pos = specimen_quiescent_coeff*pneumatic-specimen_quescent_coeff * welder_spring_constant*overall_pos - welder_quiescent_coeff*pneumatic + welder_quiescent_coeff*welder_spring_constant*overall_pos
    # overall_pos*(1.0 + specimen_quiescent_coeff*welder_spring_constant - welder_quiescent_coeff*welder_spring_constant) = (specimen_quiescent_coeff - welder_quiescent_coeff)*pneumatic
    # overall_pos = (specimen_quiescent_coeff - welder_quiescent_coeff)*pneumatic/(1.0 + specimen_quiescent_coeff*welder_spring_constant - welder_quiescent_coeff*welder_spring_constant)

    # Hertzian equilibrium displacement:
    # Displacement = (9/(16Estar^2R))^(1/3) * contact_F^(2/3)
    initial_contact_displacement = (9.0/(16.0*R_contact*Estar**2.0))**(1.0/3.0) * pneumatic_force**(2.0/3.0)

    welder_overall_pos = (convolution_evaluation.quiescent_value(1.0,specimen_resp.A,specimen_resp.alpha,specimen_resp.dt,specimen_resp.h).real - convolution_evaluation.quiescent_value(1.0,welder_tip_tip_resp_local.A,welder_tip_tip_resp_local.alpha,welder_tip_tip_resp_local.dt,welder_tip_tip_resp_local.h).real)*pneumatic_force/(1.0 + convolution_evaluation.quiescent_value(1.0,specimen_resp.A,specimen_resp.alpha,specimen_resp.dt,specimen_resp.h).real*welder_spring_constant - convolution_evaluation.quiescent_value(1.0,welder_tip_tip_resp_local.A,welder_tip_tip_resp_local.alpha,welder_tip_tip_resp_local.dt,welder_tip_tip_resp_local.h).real*welder_spring_constant) + initial_contact_displacement

    # Define convolutions of the various responses
    # This is for zero initial displacement
    #specimen_conv = convolution_evaluation.blank_from_imp_resp(specimen_resp)
    # This is for equilbrium initial displacement

    specimen_conv = convolution_evaluation.quiescent_from_imp_resp(specimen_resp,pneumatic_force-welder_overall_pos*welder_spring_constant,gpu_context_device_queue=gpu_context_device_queue,gpu_precision=gpu_precision)

    # specimen_dict_conv gets convolution evaluations setup just like
    # specimen_conv for all the other characteristics of interest
    # (laser point velocity, crack normal stress, etc.)
    specimen_dict_conv=collections.OrderedDict()
    specimen_dict_history=collections.OrderedDict()
    for specimen_motion_name in specimen_dict:
        original_unit_str = specimen_units_dict[specimen_motion_name]
        original_units = parseunits(original_unit_str)

        # Convolving with force as a function multiplies by Newton seconds
        convolve_units = parseunits("N*s")
        convolution_units = multiplyunits(original_units,convolve_units)

        specimen_motion_characteristic="%s(%s)" % (specimen_motion_name,printunits(convolution_units))

        specimen_dict_conv[specimen_motion_characteristic] = convolution_evaluation.quiescent_from_imp_resp(specimen_dict[specimen_motion_name],pneumatic_force-welder_overall_pos*welder_spring_constant,gpu_context_device_queue=gpu_context_device_queue,gpu_precision=gpu_precision)

        specimen_dict_history[specimen_motion_characteristic]=np.zeros(trange.shape[0],dtype='d')

        pass

    assert(dt==specimen_resp.dt)

    # This is for zero initial displacement
    #welder_tip_tip_conv = convolution_evaluation.blank_from_imp_resp(welder_tip_tip_resp_local)
    # This is for equilbrium initial displacement
    welder_tip_tip_conv = convolution_evaluation.quiescent_from_imp_resp(welder_tip_tip_resp_local,pneumatic_force)
    assert(dt==welder_tip_tip_resp_local.dt)

    welder_tip_elec_conv = convolution_evaluation.blank_from_imp_resp(welder_tip_elec_resp_local)
    assert(dt==welder_tip_elec_resp_local.dt)

    welder_elec_tip_conv = convolution_evaluation.blank_from_imp_resp(welder_elec_tip_resp_local)
    assert(dt==welder_elec_tip_resp_local.dt)
    assert(welder_elec_tip_resp_local.h[0]==0.0) # response of tip to electrical excitation MUST be delayed

    welder_elec_elec_conv = convolution_evaluation.blank_from_imp_resp(welder_elec_elec_resp_local)
    assert(dt==welder_elec_elec_resp_local.dt)
    assert(welder_elec_elec_resp_local.h[0]==0.0) # electrical response to electrical excitation MUST be delayed

    specimen_z_history=np.zeros(trange.shape[0],dtype='d')
    welder_tip_z_history=np.zeros(trange.shape[0],dtype='d')
    contact_F_history=np.zeros(trange.shape[0],dtype='d')
    welder_overall_velocity_history=np.zeros(trange.shape[0],dtype='d')

    welder_tip_z=welder_tip_tip_conv.evaluate() + welder_elec_tip_conv.evaluate() + welder_overall_pos

    specimen_z=specimen_conv.evaluate()

    last_overlap = welder_tip_z - specimen_z

    for tcnt in range(trange.shape[0]):
        if tcnt % 10000 == 0:
            print("tcnt=%d/%d" % (tcnt,trange.shape[0]))

            # memory leak debugging
            #import tracemalloc
            #snapshot=tracemalloc.take_snapshot()
            #top_stats=snapshot.statistics('lineno')
            #print("Top 25:")
            #print("\n".join([str(stat) for stat in top_stats[:25]]))
            #print(" ")
            pass

        #if tcnt==42890:
        #    raise ValueError("Debug!")

        #if tcnt==3953:
        #    import pdb
        #    pdb.set_trace()
        #    pass

        specimen_z=specimen_conv.step_without_instantaneous()

        welder_overall_pos += welder_overall_velocity*dt
        welder_tip_z=welder_tip_tip_conv.step_without_instantaneous() + welder_elec_tip_conv.step_without_instantaneous() + welder_overall_pos

        welder_elec_resp_voltage = welder_tip_elec_conv.step_without_instantaneous() + welder_elec_elec_conv.step_without_instantaneous()

        # Determine electrical control input

        # Welder elec input  --- Could include welder controller behavior here (based on welder_elec_resp_voltage and its history)

        if trange[tcnt] >= t0_t1 and trange[tcnt] <= t2_t3:
            welder_elec_input = welder_elec_ampl*np.cos(2*np.pi*welder_elec_freq*(trange[tcnt]-t0_t1))
            pass
        else:
            welder_elec_input = 0.0
            pass

        # Determine contact force from overlap

        # evaluate overlap
        overlap = welder_tip_z - specimen_z  # overlap represents amount of overalp between welder and specimen if no force is applied in this step

        if overlap > 0:
            # Conditions
            # Define Fwelder positive compression into welder
            # Define Fspecimen positive compression into specimen
            # Fwelder = Fspecimen
            # zshift_welder = instantaneous_welder_displacement*Fwelder*dt    # instantaneous_welder_displacement from welder model corresponds to displacement resulting from a 1 N*s impulse. Therefore it can be interpreted as have units of meters/(N*s) It is negative because the positive contact_F pulse pushes the welder away from the vibrometer
            # zshift_specimen = specimen_response[0]*Fspecimen*dt # Due to 1 N*s impulse... Should be positive
            # Add in contact springiness in series with
            # surface springiness:
            #  zspring=contact_F/kspring
            # zshift_welder - zshift_specimen - zspring = -overlap (known)

            # Solve this sytem....
            # let contact_F = Fwelder = Fspecimen = Fspring
            # instantaneous_welder_displacement*contact_F*dt  - specimen_response[0]*contact_F*dt - contact_F/kspring = -overlap
            # F = -overlap/(dt*(instantaneous_welder_displacement-specimen_resp[0]) - (1/kspring))

            # New Hertzian contact model
            # instantaneous_welder_displacement*contact_F*dt  - specimen_response[0]*contact_F*dt - (9/(16Estar^2R))^(1/3) * contact_F^(2/3) = -nocontactforce_overlap

            # This is a cubic equation for contact_F
            # for a*F - b*F^(2/3) + c = 0
            # G = F^(1/3)
            # a*G^3 - b*G^2 + c = 0
            G = np.roots([np.real(welder_tip_tip_resp_local.h[0]*dt-specimen_resp.h[0]*dt),-(9.0/(16.0*R_contact*Estar**2.0))**(1.0/3.0),0.0,overlap])
            use_G = G[(G.real > 0) & (abs(G.imag) <= G_nominal*1e-8)].real
            if use_G.shape[0] != 1:
                raise ValueError("Bad Hertzian contact solution: %s" % (str(G)))

            contact_F=use_G[0]**3.0

            # Old simple spring contact model here:
            #contact_F = -overlap/(dt*(np.real(welder_tip_tip_resp_local.h[0]-specimen_resp.h[0])) - 1.0/kspring)  # welder_tip_tip_resp_local.h[0] is negative... specimen_resp.h[0] is positive, overlap is positive, so contact_F is positive for compressive force between
            # Welder and specimen
            # Accumulate immediate response to this force
            pass
        else:
            contact_F=0.0
            pass

        contact_F_history[tcnt]=contact_F

        # Need to convolve force history
        # displacement = integral( F(t-tau) * response(tau)) dtau
        # store force history Fhist(t)
        # At a particular time t,
        # displacement = integral( F(t-tau) * response(tau)) dtau where tau >= 0

        #welder_tip_z = welder_tip_z + contact_F*welder_tip_tip_resp_local.h[0]*dt
        #specimen_z = specimen_z + contact_F*specimen_resp.h[0]*dt

        # Positive (compressive) contact_F gives positive instantaneous
        # contribution... specimen_z moves in +z direction (away from welder)
        specimen_z += specimen_conv.step_instantaneous_contribution(contact_F)

        for specimen_motion_characteristic in specimen_dict_conv:
            specimen_dict_history[specimen_motion_characteristic][tcnt] = specimen_dict_conv[specimen_motion_characteristic].step(contact_F)
            pass

        # Positive (compressive) contact_F gives negative instantaneous
        # contribution... welder tip moves in -z direction (away from specimen)
        # Also need to apply contribution of welder electric input (but
        # the instantaneous effect of this is zero
        welder_tip_z += welder_tip_tip_conv.step_instantaneous_contribution(contact_F) + welder_elec_tip_conv.step_instantaneous_contribution(welder_elec_input)

        # Apply the contribution to the electric response...
        # instantaneous effect is zero.
        welder_elec_resp_voltage = welder_tip_elec_conv.step_instantaneous_contribution(contact_F) + welder_elec_elec_conv.step_instantaneous_contribution(welder_elec_input)
        assert(welder_elec_resp_voltage==0.0) # These should not have instantaneous responses!

        # Should verify that overlap has been reduced to approximately 0
        new_overlap = welder_tip_z - specimen_z
        #assert(new_overlap < 1e-12)  (new overlap is no longer small now that we have added contact stiffness)

        specimen_z_history[tcnt]=specimen_z
        welder_tip_z_history[tcnt]=welder_tip_z

        # determine change in welder_overall_velocity
        # ***!!!! Should this have a lag ?
        # F = ma  -> a = F/m
        # v = vprior + a*dt
        # v = vprior + (F/m)*dt
        # Where F = pneumatic_force - contact_F - welder_overall_dashpot*welder_overall_velocity
        welder_spring_force = -welder_tip_z*welder_spring_constant + pneumatic_force    # use "pneumatic_force" parameter as spring preload
        welder_overall_velocity += ((welder_spring_force-contact_F - welder_overall_dashpot*welder_overall_velocity)/mass_of_welder_and_slider) * dt
        assert(np.imag(welder_overall_velocity)==0.0)

        welder_overall_velocity_history[tcnt]=welder_overall_velocity

        #if tcnt >= 80890 and contact_F > 0:
        #raise ValueError("Debug!")
        #    break
        last_overlap = new_overlap
        pass
    # !!! Need conservation of energy constraint !!!***
    # Would it help to make contact more compliant (don't require
    #  exactly zero overlap)?

    motiontable = pd.DataFrame(index=pd.Float64Index(data=np.arange(trange.shape[0],dtype='d')*dt,dtype='d',name="Time(s)"))

    motiontable.insert(len(motiontable.columns),"specimen_z_history(m)",specimen_z_history)
    motiontable.insert(len(motiontable.columns),"welder_tip_z_history(m)",welder_tip_z_history)
    motiontable.insert(len(motiontable.columns),"contact_F_history(N)",contact_F_history)
    motiontable.insert(len(motiontable.columns),"welder_overall_velocity_history(m/s)",welder_overall_velocity_history)

    motiontable.insert(len(motiontable.columns),"welder_tip_tip_resp(m/(N*s))",welder_tip_tip_resp_local.eval(np.arange(trange.shape[0])).real)
    motiontable.insert(len(motiontable.columns),"specimen_resp(m/(N*s))",specimen_resp.eval(np.arange(trange.shape[0])).real)

    for specimen_motion_characteristic in specimen_dict_history:
        motiontable.insert(len(motiontable.columns),specimen_motion_characteristic,specimen_dict_history[specimen_motion_characteristic])
        pass

    return motiontable

def write_motiontable(motiontable,output_filename):
    if output_filename.endswith(".bz2"):
        motiontable.to_csv(output_filename,compression='bz2')
        pass
    else:
        motiontable.to_csv(output_filename)
        pass

    pass

def plot_contact(motiontable,exc_t0):

    trange = np.array(motiontable.index) # ["Time(s)"]
    dt=trange[1]-trange[0]

    difft = (trange[:-1]+trange[1:])/2.0

    max_t_plot = np.max(trange)

    # round length up to next power-of-two size
    # so we don't get bogged down if trange.shape[0] is prime
    fft_len = 1 << int(np.ceil(np.log(trange.shape[0])/np.log(2)))

    frange = np.arange(fft_len,dtype='d')/(fft_len*dt)
    frange[fft_len//2:] -= 1.0/dt
    frange[fft_len//2]=np.nan

    welder_tip_z_history = motiontable["welder_tip_z_history(m)"]
    specimen_z_history = motiontable["specimen_z_history(m)"]
    contact_F_history = motiontable["contact_F_history(N)"]
    specimen_laser_vel_history = motiontable["specimen_laser(m/s)"]

    impresp_plot=pl.figure()
    pl.clf()
    pl.title("Welder and specimen impulse response")
    pl.plot(trange*1e6,
            motiontable["welder_tip_tip_resp(m/(N*s))"],'-',
            trange*1e6,
            motiontable["specimen_resp(m/(N*s))"],'--')
    pl.xlabel('time (us)')
    pl.legend(('Welder','Specimen'))
    pl.grid()

    velspec_plot=pl.figure()
    pl.clf()
    pl.title("Welder and specimen velocity spectrum")
    pl.plot(frange/1e3,
            np.abs(np.fft.fft(motiontable["welder_tip_tip_resp(m/(N*s))"],n=fft_len)*dt*(2.0*np.pi*np.abs(frange))),'-',
            frange/1e3,
            np.abs(np.fft.fft(motiontable["specimen_resp(m/(N*s))"],n=fft_len)*dt*(2.0*np.pi*np.abs(frange))),'--')
    pl.xlabel('Frequency (kHz)')
    pl.ylabel('Velocity spectrum (m/s/Hz)')
    pl.legend(('Welder','Specimen'))
    pl.axis((0,250,0,.012))
    pl.grid()

    phasespec_plot = pl.figure()
    # NOTE: Because welder velocity is negated
    # relative to force direction, we undo that by negating before
    # evaluating the angle, so it is clear whether the phase is in the
    # required (-pi/2,pi/2) range
    pl.clf()
    pl.title("Welder and specimen velocity phase spectrum")
    pl.plot(frange/1e3,
            np.angle(-np.fft.fft(motiontable["welder_tip_tip_resp(m/(N*s))"],n=fft_len)*dt*((0+1j)*2.0*np.pi*np.abs(frange))),'-',
            frange/1e3,
            np.angle(np.fft.fft(motiontable["specimen_resp(m/(N*s))"],n=fft_len)*dt*((0+1j)*2.0*np.pi*np.abs(frange))),'--')
    pl.xlabel('Frequency (kHz)')
    pl.ylabel('Phase angle (rad)')
    pl.legend(('Welder','Specimen'))
    pl.axis((0,250,-np.pi,np.pi))
    pl.grid()

    disp_plot = pl.figure()
    pl.clf()
    pl.plot(trange*1e3,specimen_z_history*1e6,'-',
            trange*1.e3,welder_tip_z_history*1e6,'-')
    pl.axis((0,max_t_plot*1.e3,min(np.min(specimen_z_history*1e6),np.min(welder_tip_z_history*1e6)),max(np.max(specimen_z_history*1e6),np.max(welder_tip_z_history*1e6))))
    pl.xlabel('Time (ms)')
    pl.ylabel('Displacement (um)')
    pl.title('Specimen and welder displacement')
    pl.legend(('Specimen','Welder'))
    pl.grid()

    dispzoom_plot = pl.figure()
    pl.clf()
    fig5_tstart = exc_t0 - 0.1e-3 #5e-3
    fig5_tend = exc_t0 + 0.9e-3 # 6e-3
    pl.plot(trange*1e3,specimen_z_history*1e6,'-',
            trange*1.e3,welder_tip_z_history*1e6,'-')
    pl.axis((fig5_tstart*1e3,fig5_tend*1e3,min(np.min(specimen_z_history[(trange >= fig5_tstart) & (trange <= fig5_tend)]*1e6),np.min(welder_tip_z_history[(trange >= fig5_tstart) & (trange <= fig5_tend)]*1e6)),max(np.max(specimen_z_history[(trange >= fig5_tstart) & (trange <= fig5_tend)]*1e6),np.max(welder_tip_z_history[(trange >= fig5_tstart) & (trange <= fig5_tend)]*1e6))))
    pl.xlabel('Time (ms)')
    pl.ylabel('Displacement (um)')
    pl.title('Specimen and welder displacement')
    pl.legend(('Specimen','Welder'))
    pl.grid()

    contactforce_plot = pl.figure()
    pl.clf()
    # Impulse between .28*.33
    Impulse=contact_F_history[(trange >.28e-3) & (trange < .33e-3)].sum()*dt
    pl.plot(trange*1e3,contact_F_history)
    pl.axis((0,max_t_plot*1.e3,0,20000))
    pl.xlabel('Time (ms)')
    pl.ylabel('Force (N)')
    pl.title('Contact force (time domain')
    pl.grid()

    #pl.figure(5)
    #pl.clf()
    #pl.plot(trange,welder_overall_velocity_history)

    contactspectrum_plot = pl.figure()
    pl.clf()
    # Impulse between .28*.33
    pl.plot(frange/1e3,np.abs(np.fft.fft(contact_F_history,n=fft_len)*dt))
    #pl.axis((0,max_t_plot*1.e3,0,20000))
    pl.xlabel('Frequency (kHz)')
    pl.ylabel('Force spectrum (N/Hz)')
    pl.title('Contact force (frequency domain')
    pl.grid()

    overlap_plot = pl.figure()
    pl.clf()
    pl.plot(trange*1.e3,(welder_tip_z_history-specimen_z_history)*1e6,'-')
    #pl.axis((0,55,-1000,500))
    pl.xlabel('Time (ms)')
    pl.ylabel('Displacement (um)')
    pl.title('Welder/specimen overlap')
    #pl.axis((0,1,-3,1))
    pl.grid()

    #pl.figure(7)
    #pl.clf()
    #difft = (trange[:-1]+trange[1:])/2.0
    #pl.plot(difft*1.e3,np.diff(welder_tip_z_history-specimen_z_history)/dt,'-')
    #pl.xlabel('Time (ms)')
    #pl.ylabel('Velocity (m/s)')
    #pl.grid()

    contactvel_plot = pl.figure()
    pl.clf()
    pl.plot(difft*1.e3,np.diff(specimen_z_history)/dt,'-')
    #pl.axis((0,55,-1000,500))
    pl.xlabel('Time (ms)')
    pl.ylabel('Velocity (m/s)')
    pl.title('Specimen contact surface velocity (synthetic vibrometer)')
    #pl.axis((0,1,-3,1))
    pl.grid()

    contactvelspec_plot = pl.figure()
    pl.clf()
    pl.plot(frange/1.e3,np.abs(np.fft.fft(specimen_z_history,n=fft_len)*dt*(2.0*np.pi*frange)),'-')
    #pl.axis((0,55,-1000,500))
    pl.xlabel('Frequency (kHz)')
    pl.ylabel('Velocity spectrum (m/s)/Hz')
    pl.title('Specimen contact velocity spectrum (synthetic vibrometer)')
    #pl.axis((0,1,-3,1))
    pl.grid()

    laservel_plot = pl.figure()
    pl.clf()
    pl.plot(trange*1e3,specimen_laser_vel_history,'-')
    pl.xlabel('Time (ms)')
    pl.ylabel('Velocity (m/s)')
    pl.title('Laser spot surface velocity (synthetic vibrometer)')
    #pl.axis((0,1,-3,1))
    pl.grid()

    laservelspec_plot = pl.figure()
    pl.clf()
    pl.plot(frange/1.e3,np.abs(np.fft.fft(specimen_laser_vel_history,n=fft_len)*dt),'-')
    #pl.axis((0,55,-1000,500))
    pl.xlabel('Frequency (kHz)')
    pl.ylabel('Velocity spectrum (m/s)/Hz')
    pl.title('Laser spot surface velocity spectrum (synthetic vibrometer)')
    #pl.axis((0,1,-3,1))
    pl.grid()

    plotdict = {
        "impulse_response": impresp_plot,
        "velocity_spectrum": velspec_plot,
        "phase_spectrum": phasespec_plot,
        "displacement": disp_plot,
        "displacement_zoom": dispzoom_plot,
        "contact_force": contactforce_plot,
        "contact_spectrum": contactspectrum_plot,
        "overlap": overlap_plot,
        "contact_velocity": contactvel_plot,
        "contact_velocity_spectrum": contactvelspec_plot,
        "laser_velocity": laservel_plot,
        "laser_velocity_spectrum": laservelspec_plot,
    }
    return plotdict


import abc
from typing import Callable, Optional, Tuple, TypeVar

import equinox as eqx
import jax
import jax.numpy as jnp

from ..custom_types import Bool, DenseInfo, PyTree, PyTreeDef, Scalar
from ..heuristics import is_sde
from ..local_interpolation import AbstractLocalInterpolation
from ..nonlinear_solver import AbstractNonlinearSolver, NewtonNonlinearSolver
from ..solution import RESULTS
from ..term import AbstractTerm

_SolverState = TypeVar("SolverState", bound=Optional[PyTree])

def vector_tree_dot(a, b):
    return jax.tree_map(lambda bi: jnp.tensordot(a, bi, axes=1), b)

class _MetaAbstractSolver(type(eqx.Module)):
    def __instancecheck__(cls, obj):
        if super(_MetaAbstractSolver, AbstractWrappedSolver).__instancecheck__(obj):
            obj = obj.solver
        return super().__instancecheck__(obj)

class AbstractSolver(eqx.Module, metaclass=_MetaAbstractSolver):
    """Abstract base class for all differential equation solvers."""

    @property
    @abc.abstractmethod
    def term_structure(self) -> PyTreeDef:
        """What PyTree structure `terms` should have when used with this solver."""

    # On the type: frequently just Type[AbstractLocalInterpolation]
    @property
    @abc.abstractmethod
    def interpolation_cls(self) -> Callable[..., AbstractLocalInterpolation]:
        """How to interpolate the solution in between steps."""

    def order(self, terms: PyTree[AbstractTerm]) -> Optional[int]:
        """Order of the solver for solving ODEs."""
        return None

    def strong_order(self, terms: PyTree[AbstractTerm]) -> Optional[Scalar]:
        """Strong order of the solver for solving SDEs."""
        return None

    def error_order(self, terms: PyTree[AbstractTerm]) -> Optional[Scalar]:
        """Order of the error estimate used for adaptive stepping.

        The default (slightly heuristic) implementation is as follows.

        The error estimate is assumed to come from the difference of two methods. If
        these two methods have orders `p` and `q` then the local order of the error
        estimate is `min(p, q) + 1` for an ODE and `min(p, q) + 0.5` for an SDE.

        - In the SDE case then we assume `p == q == solver.strong_order()`.
        - In the ODE case then we assume `p == q + 1 == solver.order()`.
        - We assume that non-SDE/ODE cases do not arise.

        This is imperfect as these assumptions may not be true. In addition in the SDE
        case, then solvers will sometimes exhibit higher orders of convergence for
        specific noise types (see issue #47).
        """
        if is_sde(terms):
            order = self.strong_order(terms)
            if order is not None:
                order = order + 0.5
            return order
        else:
            return self.order(terms)

    def init(
        self,
        terms: PyTree[AbstractTerm],
        t0: Scalar,
        t1: Scalar,
        y0: PyTree,
        args: PyTree,
    ) -> _SolverState:
        """Initialises any hidden state for the solver.

        **Arguments** as [`diffrax.diffeqsolve`][].

        **Returns:**

        The initial solver state, which should be used the first time `step` is called.
        """
        return None

    @abc.abstractmethod
    def step(
        self,
        terms: PyTree[AbstractTerm],
        t0: Scalar,
        t1: Scalar,
        y0: PyTree,
        args: PyTree,
        solver_state: _SolverState,
        made_jump: Bool,
    ) -> Tuple[PyTree, Optional[PyTree], DenseInfo, _SolverState, RESULTS]:
        """Make a single step of the solver.

        Each step is made over the specified interval $[t_0, t_1]$.

        **Arguments:**

        - `terms`: The PyTree of terms representing the vector fields and controls.
        - `t0`: The start of the interval that the step is made over.
        - `t1`: The end of the interval that the step is made over.
        - `y0`: The current value of the solution at `t0`.
        - `args`: Any extra arguments passed to the vector field.
        - `solver_state`: Any evolving state for the solver itself, at `t0`.
        - `made_jump`: Whether there was a discontinuity in the vector field at `t0`.
            Some solvers (notably FSAL Runge--Kutta solvers) usually assume that there
            are no jumps and for efficiency re-use information between steps; this
            indicates that a jump has just occurred and this assumption is not true.

        **Returns:**

        A tuple of several objects:

        - The value of the solution at `t1`.
        - A local error estimate made during the step. (Used by adaptive step size
            controllers to change the step size.) May be `None` if no estimate was
            made.
        - Some dictionary of information that is passed to the solver's interpolation
            routine to calculate dense output. (Used with `SaveAt(ts=...)` or
            `SaveAt(dense=...)`.)
        - The value of the solver state at `t1`.
        - An integer (corresponding to `diffrax.RESULTS`) indicating whether the step
            happened successfully, or if (unusually) it failed for some reason.
        """

    def func_for_init(
        self, terms: PyTree[AbstractTerm], t0: Scalar, y0: PyTree, args: PyTree
    ) -> PyTree:
        """Provides vector field evaluations to select the initial step size.

        This is used to make a point evaluation. This is unlike
        [`diffrax.AbstractSolver.step`][], which operates over an interval.

        In general differential equation solvers are interval-based. There is precisely
        one place where point evaluations are needed: selecting the initial step size
        automatically in an ODE solve. And that is what this function is for.

        **Arguments:** As [`diffrax.diffeqsolve`][]

        **Returns:**

        The evaluation of the vector field at `t0`.
        """

        raise ValueError(
            "An initial step size cannot be selected automatically. The most common "
            "scenario for this error to occur is when trying to use adaptive step "
            "size solvers with SDEs. Please specify an initial `dt0` instead."
        )

class AbstractImplicitSolver(AbstractSolver):
    nonlinear_solver: AbstractNonlinearSolver = NewtonNonlinearSolver()

class AbstractItoSolver(AbstractSolver):
    pass

class AbstractStratonovichSolver(AbstractSolver):
    pass

class AbstractAdaptiveSolver(AbstractSolver):
    pass

class AbstractAdaptiveSDESolver(AbstractAdaptiveSolver):
    pass

class AbstractWrappedSolver(AbstractSolver):
    solver: AbstractSolver

class HalfSolver(AbstractWrappedSolver, AbstractAdaptiveSDESolver):
    """Wraps another solver, trading cost in order to provide error estimates. (These
    error estimates mean that the solver can be used with an adaptive step size
    controller, like [`diffrax.PIDController`][].)

    For every step of the wrapped solver, it does this by also making two half-steps,
    and comparing the results. (Hence the name "HalfSolver".)

    As such each step costs 3 times the computational cost of the wrapped solver.

    !!! tip

        Many solvers already provided error estimates, making `HalfSolver` primarily
        useful when using a solver that doesn't provide error estimates -- e.g.
        [`diffrax.Euler`][] -- in particular this is common when solving SDEs.
    """

    @property
    def term_structure(self):
        return self.solver.term_structure

    @property
    def interpolation_cls(self):
        return self.solver.interpolation_cls

    def order(self, terms: PyTree[AbstractTerm]) -> Optional[int]:
        return self.solver.order(terms)

    def strong_order(self, terms: PyTree[AbstractTerm]) -> Optional[Scalar]:
        return self.solver.strong_order(terms)

    def error_order(self, terms: PyTree[AbstractTerm]) -> Optional[Scalar]:
        if is_sde(terms):
            order = self.strong_order(terms)
            if order is not None:
                order = order + 0.5
        else:
            order = self.order(terms)
            if order is not None:
                order = order + 1
        return order

    def init(
        self,
        terms: PyTree[AbstractTerm],
        t0: Scalar,
        t1: Scalar,
        y0: PyTree,
        args: PyTree,
    ):
        return self.solver.init(terms, t0, t1, y0, args)

    def step(
        self,
        terms: PyTree[AbstractTerm],
        t0: Scalar,
        t1: Scalar,
        y0: PyTree,
        args: PyTree,
        solver_state: _SolverState,
        made_jump: Bool,
    ) -> Tuple[PyTree, Optional[PyTree], DenseInfo, _SolverState, RESULTS]:

        original_solver_state = solver_state
        thalf = t0 + 0.5 * (t1 - t0)

        yhalf, _, _, solver_state, result1 = self.solver.step(
            terms, t0, thalf, y0, args, solver_state, made_jump
        )
        y1, _, _, solver_state, result2 = self.solver.step(
            terms, thalf, t1, yhalf, args, solver_state, made_jump=False
        )

        # TODO: use dense_info from the pair of half-steps instead
        y1_alt, _, dense_info, _, result3 = self.solver.step(
            terms, t0, t1, y0, args, original_solver_state, made_jump
        )

        y_error = jnp.abs(y1 - y1_alt)
        result = jnp.maximum(result1, jnp.maximum(result2, result3))

        return y1, y_error, dense_info, solver_state, result

    def func_for_init(
        self, terms: PyTree[AbstractTerm], t0: Scalar, y0: PyTree, args: PyTree
    ):
        return self.solver.func_for_init(terms, t0, y0, args)

HalfSolver.__init__.__doc__ = """**Arguments:**

- `solver`: The solver to wrap.
"""


"""Calculates the Frechet Inception Distance (FID) to evalulate GANs

The FID metric calculates the distance between two distributions of images.
Typically, we have summary statistics (mean & covariance matrix) of one
of these distributions, while the 2nd distribution is given by a GAN.

When run as a stand-alone program, it compares the distribution of
images that are stored as PNG/JPEG at a specified location with a
distribution given by summary statistics (in pickle format).

The FID is calculated by assuming that X_1 and X_2 are the activations of
the pool_3 layer of the inception net for generated samples and real world
samples respectively.

See --help to see further details.

Code apapted from https://github.com/bioinf-jku/TTUR to use PyTorch instead
of Tensorflow

Copyright 2018 Institute of Bioinformatics, JKU Linz

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.
"""

import numpy as np
import torch
from scipy import linalg
from torch.nn.functional import adaptive_avg_pool2d

from models.inception import InceptionV3

class FIDScoreCumulative:
    def __init__(self, dims=2048, device='cpu'):
        """Init FIDScore

        Params:
        -- dims        : Dimensionality of features returned by Inception
        -- device      : Device to run calculations
        """
        self.dims = dims
        self.device = device

        block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[self.dims]
        self.model = InceptionV3([block_idx]).to(self.device)
        self.model.eval()

        self.sessions = {}

    class Session:
        def __init__(self, data_len, dims):
            self.data_len = data_len
            self.dims = dims
            self.reset_activations()

        def reset_activations(self):
            self.pred_arr = np.empty((self.data_len, self.dims))
            self.start_idx = 0

        def add_activation(self, pred):
            self.pred_arr[self.start_idx:self.start_idx + pred.shape[0]] = pred
            self.start_idx = self.start_idx + pred.shape[0]

    def add_session(self, id, data_len):
        self.sessions[id] = self.Session(data_len, self.dims)

    def remove_session(self, id):
        del self.sessions[id]

    def reset_activations(self, id):
        self.sessions[id].reset_activations()

    def add_activation(self, id, batch, inf_model=None):
        """Calculates the activations of the pool_3 layer for all images.

        Params:
        -- dataloader  : Dataloader containing images or inputs to inf_model
        -- inf_model   : Model in which to input the dataset. Use dataset directly
                         if inf_model is None.

        Returns:
        -- A numpy array of dimension (num images, dims) that contains the
           activations of the given tensor when feeding inception with the
           query tensor.
        """

        batch = batch.to(self.device)
        pred = self.calculate_activation(batch, inf_model)
        self.sessions[id].add_activation(pred)

    def calculate_activation(self, batch, inf_model=None):
        # Use output of inf_model instead of batch if an inf_model is given.
        with torch.no_grad():
            if inf_model is not None:
                batch = inf_model(batch)
            pred = self.model(batch)[0]

        # If model output is not scalar, apply global spatial average pooling.
        # This happens if you choose a dimensionality not equal 2048.
        if pred.size(2) != 1 or pred.size(3) != 1:
            pred = adaptive_avg_pool2d(pred, output_size=(1, 1))

        pred = pred.squeeze(3).squeeze(2).cpu().numpy()

        return pred

    def calculate_activation_statistics(self, id):
        """Calculation of the statistics used by the FID.
        Params:
        -- dataloader  : Dataloader containing images or inputs to inf_model
        -- inf_model   : Model in which to input the dataset. Use dataset directly
                         if inf_model is None.

        Returns:
        -- mu    : The mean over samples of the activations of the pool_3 layer of
                   the inception model.
        -- sigma : The covariance matrix of the activations of the pool_3 layer of
                   the inception model.
        """
        mu = np.mean(self.sessions[id].pred_arr, axis=0)
        sigma = np.cov(self.sessions[id].pred_arr, rowvar=False)
        return mu, sigma

    def calculate_frechet_distance(self, mu1, sigma1, mu2, sigma2, eps=1e-6):
        """Numpy implementation of the Frechet Distance.
        The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1)
        and X_2 ~ N(mu_2, C_2) is
                d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)).

        Stable version by Dougal J. Sutherland.

        Params:
        -- mu1   : Numpy array containing the activations of a layer of the
                   inception net (like returned by the function 'get_predictions')
                   for generated samples.
        -- mu2   : The sample mean over activations, precalculated on an
                   representative data set.
        -- sigma1: The covariance matrix over activations for generated samples.
        -- sigma2: The covariance matrix over activations, precalculated on an
                   representative data set.

        Returns:
        --   : The Frechet Distance.
        """

        mu1 = np.atleast_1d(mu1)
        mu2 = np.atleast_1d(mu2)

        sigma1 = np.atleast_2d(sigma1)
        sigma2 = np.atleast_2d(sigma2)

        assert mu1.shape == mu2.shape, \
            'Training and test mean vectors have different lengths'
        assert sigma1.shape == sigma2.shape, \
            'Training and test covariances have different dimensions'

        diff = mu1 - mu2

        # Product might be almost singular
        covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
        if not np.isfinite(covmean).all():
            msg = ('fid calculation produces singular product; '
                   'adding %s to diagonal of cov estimates') % eps
            print(msg)
            offset = np.eye(sigma1.shape[0]) * eps
            covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))

        # Numerical error might give slight imaginary component
        if np.iscomplexobj(covmean):
            if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
                m = np.max(np.abs(covmean.imag))
                raise ValueError('Imaginary component {}'.format(m))
            covmean = covmean.real

        tr_covmean = np.trace(covmean)

        return (diff.dot(diff) + np.trace(sigma1) +
                np.trace(sigma2) - 2 * tr_covmean)


# -.-.-.-.-.-.-.-.-.-.-.-.-.-.-.-.-.-.-.-.-.
#
#  File Name : to_tardis_mapper.py
#
#  Purpose : Generate Tardis input from generic explosion models
#
#  Creation Date : 29-02-2016
#
#  Last Modified : Fri 04 Mar 2016 15:23:22 CET
#
#  Created By : unoebauer
#
# _._._._._._._._._._._._._._._._._._._._._.
"""A simple tool to map the output of SN explosion calculations or any ejecta
model into Tardis, using its capability to work with specific density and
abundance files.
"""
import numpy as np
import logging
import astropy.units as units
from pyne import nucname, material

logger = logging.getLogger(__name__)

try:
    material.Material().decay
except AttributeError:
    logger.critical("PyNe module outdated: version >= 0.5 is required")
    raise ImportError("No recent PyNe module found")

# maximum atomic number
zmax = 30

class original_model(object):
    """Simple model interface. It and its derived classes should provide a
    common interface for all possible explosion model formats.

    This class only requires a minimal set of information, from which all the
    remaining quantities are constructed (under the assumption of homologous
    expansion). For each radial shell, the inner and the outer shell radius
    have to be provided, in addition with the shell density, the mass fractions
    of all stable elements and of all radioactive elements. The final
    information to be passed is the time since explosion.

    Velocity and shell masses are generated from these input data, by assuming
    perfect homologous expansion, (i.e. $u = r/t$) and spherical symmetry.

    Notes
    -----

    The to_tardis_mapper tool expects that stable_abundances + radio_abundance
    = 1. In other words, the radioactive isotopes should not be contained in
    the mass fractions stores in stable_abundances.

    Parameters
    ----------
    """
    def __init__(self):

        self._ro = None
        self._ri = None
        self._t = None
        self._rho = None

        self._nzones = None
        self._vo = None
        self._vi = None
        self._dm = None
        self._mr = None

        self._stable_abundances = None
        self._radio_abundances = None

    def _reset_cached_variables(self):
        """Resets the cached derived quantities. This routine should be called
        every time one of the principle inputs (such as density) is reset.
        """

        self._nzones = None
        self._vo = None
        self._vi = None
        self._dm = None
        self._mr = None

    @property
    def ro(self):
        """outer shell edge for each radial cell"""
        if self._ro is None:
            raise ValueError(
                "You have to read-in or assign the outer cell radii first")
        return self._ro

    @ro.setter
    def ro(self, val):
        try:
            val.to("cm")
        except (AttributeError, units.UnitConversionError):
            raise ValueError(
                "ro must be a valid length astropy.units.Quantity")
        self._ro = val
        self._reset_cached_variables()

    @property
    def ri(self):
        """inner shell edge for each radial cell"""
        if self._ri is None:
            raise ValueError(
                "You have to read-in or assign the inner cell radii first")
        return self._ri

    @ri.setter
    def ri(self, val):
        try:
            val.to("cm")
        except (AttributeError, units.UnitConversionError):
            raise ValueError(
                "ri must be a valid length astropy.units.Quantity")
        self._ri = val
        self._reset_cached_variables()

    @property
    def rho(self):
        """cell density"""
        if self._rho is None:
            raise ValueError(
                "You have to read-in or assign the cell density first")
        return self._rho

    @rho.setter
    def rho(self, val):
        try:
            val.to("g/cm^3")
        except (AttributeError, units.UnitConversionError):
            raise ValueError(
                "rho must be a valid mass density astropy.units.Quantity")
        self._rho = val
        self._reset_cached_variables()

    @property
    def t(self):
        """time since explosion"""
        if self._t is None:
            raise ValueError(
                "You have to read-in or assign the time since explosion first")
        return self._t

    @t.setter
    def t(self, val):
        try:
            val.to("s")
        except (AttributeError, units.UnitConversionError):
            raise ValueError(
                "t must be a valid time astropy.units.Quantity")
        self._t = val
        self._reset_cached_variables()

    @property
    def nzones(self):
        """number of shells in the model"""
        if self._nzones is None:
            self._nzones = len(self.ro)

        return self._nzones

    @property
    def vo(self):
        """fluid velocity at outer cell edge"""
        if self._vo is None:
            self._vo = self.ro / self.t
        return self._vo

    @property
    def vi(self):
        """fluid velocity at inner cell edge"""
        if self._vi is None:
            self._vi = self.ri / self.t
        return self._vi

    @property
    def dm(self):
        """mass contained within a cell"""
        if self._dm is None:
            self._dm = 4. * np.pi / 3. * (self.ro**3 - self.ri**3) * self.rho
        return self._dm

    @property
    def mr(self):
        """mass enclosed by shells outer radius"""
        if self._mr is None:
            _mr = np.zeros(self.nzones) * self.dm.unit
            _mr[0] = self.dm[0]
            for i in xrange(1, self.nzones):
                _mr[i] = _mr[i-1] + self.dm[i]
            self._mr = _mr
        return self._mr

    @property
    def stable_abundances(self):
        """a dictionary holding the mass fractions of all stable elements up to
        Z = zmax. The proton number serves as dictionary key
        """
        if self._stable_abundances is None:
            tmp = {}

            for i in xrange(zmax):
                z = i+1
                tmp[z] = np.zeros(self.nzones)
            self._stable_abundances = tmp

        return self._stable_abundances

    @property
    def radio_abundances(self):
        """a dictionary holding the mass fractions of all radioactive isotopes
        at time = t. Strings such as 'Ni56', i.e. containing the Element symbol
        and the mass number, serve as dictionary keys
        """
        if self._radio_abundances is None:
            self._radio_abundances = {}
        return self._radio_abundances

    @property
    def complete(self):
        """A simple flag determining whether all quantities necessary for the
        mapping process have been set/read-in"""

        _complete = False
        try:
            self.t
            self.rho
            self.ro
            self.ri
            _complete = True
        except ValueError:
            pass

        try:
            X = np.array(
                [self.stable_abundances[z] for z in xrange(1, zmax+1)])
            _complete = _complete * (X > 0).any()
        except ValueError:
            _complete = False

        return _complete

    def read_density(self, fname):
        """A prototype for reading the radius and density information from
        file"""

        pass

    def read_abundances(self, fname):
        """A prototype for reading the stable and radioactive elemental
        abundances from file"""

        pass

class w7_model(original_model):
    """A simple interface class for the W7 model.

    Notes
    -----
    The original W7 model has been presented by [1]_. This reader is designed
    particular version calculated by [2]_.

    References
    ----------

    .. [1] Nomoto et al. "Accreting white dwarf models of Type I supernovae.
       III - Carbon deflagration supernovae" ApJ, 1984, 286, 644-658
    .. [2] Iwamoto et al. "Nucleosynthesis in Chandrasekhar Mass Models for
       Type IA Supernovae and Constraints on Progenitor Systems and
       Burning-Front Propagation" ApJS, 1999, 125, 439-462
    """
    def __init__(self):
        super(w7_model, self).__init__()

    def read_density(self, fname):
        """Read the radius and density of the W7 model from file.

        Parameters
        ----------
        fname : str
            Name of the hydrodynamics file of the W7 model
        """

        f = open(fname, "r")

        # read header
        buffer = f.readline().rsplit()
        self.t = float(buffer[2]) * units.s
        buffer = f.readline()

        # read main data block
        data = np.loadtxt(f)
        f.close()

        self.ro = data[:, 2] * units.cm
        self.ri = np.insert(self.ro, 0, 0 * units.cm)[:-1]
        self.rho = data[:, 3] * units.g / units.cm**3

    def read_abundances(self, fname):
        """Read elemental abundances of the W7 model from file.

        Parameters
        ----------
        fname : str
            Name of the nucleosynthesis file of the W7 model
        """

        f = open(fname, "r")
        data = np.loadtxt(f, skiprows=1)
        f.close()

        # the file contains abundances for elements from Z=1 to Z=32 (Ge).
        for i in xrange(np.max([zmax, 32])):
            self.stable_abundances[i+1] = data[:, i]

        # the last three columns contain the mass fractions of the radioactive
        # isotopes nickel-56, cobalt-56, nickel-57
        self.radio_abundances["ni56"] = data[:, 32]
        self.radio_abundances["co56"] = data[:, 33]
        self.radio_abundances["ni57"] = data[:, 34]

class to_tardis_mapper(object):
    """A simple mapper object, transforming explosion models, stored in the
    original_model interface classes, into Tardis specific structure input
    files

    Parameters
    ----------
    orig_model : original_model or derived classes
        interface object holding the essential data of the original explosion
        model
    """
    def __init__(self, orig_model):

        if not orig_model.complete:
            logger.critical("Not all necessary information have been"
                            " set/read-in in the original model")
            raise ValueError("Check original_model for completeness")
        self.orig = orig_model

    def remap(self, v, t, decay=True, write_density=True,
              density_fname="tardis_densities.dat", write_abundances=True,
              abundance_fname="tardis_abundances.dat", be_fix=True, be_to_z=6):
        """Perform the remapping of the original model onto a specified Tardis
        velocity grid. A homologous expansion from the time of the model to the
        time since explosion used in the Tardis calculation is automatically
        performed. Optionally, the decay of the radioactive isotopes can also
        be performed. The remapped model can then be written to Tardis files.

        Notes
        -----
        When supplying specific structure files to Tardis, it expects to find
        also the density and composition of the photosphere. Thus, if the
        ejecta above the photosphere is supposed to be described by 20 shells,
        the data files have to contain 21 rows. The specific density and
        composition of the photosphere is not important since this information
        is not used within Tardis - the photospheric velocity (i.e. the
        velocity of the first row) is used!

        Parameters
        ----------
        v : array-like astropy.units.Quantity
            Tardis velocity grid, interpreted as the velocity at the outer
            shell edges. Mind the photospheric shell!
        t : scalar astropy.units.Quantity
            start time (time since explosion) of the Tardis calculation
        decay : bool
            perform decay of the radioactive isotopes (default True)
        write_density : bool
            write remapped density to a Tardis specific structure file (default
            True)
        density_fname : str
            name of the Tardis specific structure file (default
            'tardis_densities.dat')
        write_abundances : bool
            write remapped abundances to a Tardis specific abundance file
            (default True)
        abundance_fname : str
            name of the Tardis specific structure file (default
            'tardis_abundances.dat')
        be_fix : bool
            perform the Beryllium fix (see #438), i.e. set Be to zero and add
            its abundance to a different element (default True)
        be_to_z : int
            proton number of the destination element for the Be fix (default 6,
            i.e. carbon)
        """

        self._remap_density(v, t)
        self._remap_abundances()

        if decay:
            self._decay_abundances()
        else:
            self._copy_radio_abundances()
        if be_fix:
            self._be_fix(to_z=be_to_z)
        if write_density:
            self._write_tardis_density_file(fname=density_fname)
        if write_abundances:
            self._write_tardis_abundance_file(fname=abundance_fname)

    def _remap_density(self, v, t):
        """Remap density of the original_model onto the provided velocity
        grid and perform a homologous expansion of the density until the
        defined start time of Tardis

        Notes
        -----

        For the remapping, we interpolate mr in the v**3 space (corresponds to
        volume space).

        Parameters
        ----------
        v : array-like astropy.units.Quantity
            Tardis velocity grid, interpreted as the velocity at the outer
            shell edges. Mind the photospheric shell!
        t : scalar astropy.units.Quantity
            start time (time since explosion) of the Tardis calculation
        """

        self.t = t
        self.N_interp = len(v) - 1
        self.v_interp_r = v[1:].to("cm/s")
        self.v_interp_l = v[:-1].to("cm/s")

        V_interp = 4. / 3. * np.pi * ((self.v_interp_r * t).to("cm")**3 -
                                      (self.v_interp_l * t).to("cm")**3)

        vrorig = self.orig.vo.to("cm/s")
        mrorig = self.orig.mr.to("solMass")

        mr_interp = np.interp(
            np.insert(self.v_interp_r, 0, self.v_interp_l[0])**3,
            np.append(0 * vrorig.unit, vrorig)**3,
            np.append(0 * mrorig.unit, mrorig)) * mrorig.unit

        self.dm_interp = (mr_interp[1:] - mr_interp[:-1])
        self.rho_interp = (self.dm_interp / V_interp).to("g/cm^3")

    def _remap_abundances(self):
        """Remap abundances for the original model onto the grid defined in
        _remap_density. This routine has be called after _remap_density is
        performed.

        Notes
        -----
        Must be called after _remap_density.

        Raises
        ------
        AttributeError if called before _remap_density
        """
        try:
            self.v_interp_l
            self.v_interp_r
        except AttributeError:
            logger.critical("Density must be remapped before abundances are"
                            " addressed")
            raise AttributeError("no v_interp_r; call _remap_density first")

        self.abundances_interp = {}
        self.radio_abundances_interp = {}

        def remap_species(Xorig):
            """helper routine to remap one specific elemental species onto the
            new velocity grid

            Parameters
            ----------
            Xorig : numpy.ndarray
                original mass fraction of the original model

            Returns
            -------
            X_interp : numpy.ndarray
                remapped mass fractions, corresponding to the Tardis model
            """

            vrorig = self.orig.vo.to("cm/s")
            Xrorig = np.zeros(len(Xorig)) * self.orig.dm.unit

            Xrorig[0] = self.orig.dm[0] * Xorig[0]
            for i in xrange(1, self.orig.nzones):
                Xrorig[i] += Xrorig[i-1] + self.orig.dm[i] * Xorig[i]

            X_interp = np.interp(
                np.insert(self.v_interp_r, 0, self.v_interp_l[0])**3,
                np.append(0 * vrorig.unit, vrorig)**3,
                np.append(0 * Xrorig.unit, Xrorig)) * Xrorig.unit

            X_interp = X_interp[1:] - X_interp[:-1]

            return (X_interp / self.dm_interp).to("").value

        # remap stable elements
        for z in xrange(1, zmax+1):

            Xorig = self.orig.stable_abundances[z]
            X_interp = remap_species(Xorig)

            self.abundances_interp[z] = X_interp

        # remap radioactive isotopes
        for ident in self.orig.radio_abundances.keys():

            Xorig = self.orig.radio_abundances[ident]

            X_interp = remap_species(Xorig)

            self.radio_abundances_interp[ident] = X_interp

    def _copy_radio_abundances(self):
        """Copies the radioactive isotopes onto the corresponding stable
        elements. This routine is intended for uses of the mapper during which
        the decay is neglected

        Notes
        -----
        Must be called after _remap_abundances.

        Raises
        ------
        AttributeError if called before _remap_abundances
        """
        try:
            self.radio_abundances_interp
        except AttributeError:
            logger.critical("Abundances must be remapped before radioactive:"
                            " isotopes are copied onto the stable elements")
            raise AttributeError("no radio_abundances_interp; call"
                                 " _remap_abundances first")

        for ident in self.radio_abundances_interp.keys():
            elemid = nucname.id(ident)
            z = nucname.znum(elemid)
            self.abundances_interp[z] = \
                self.abundances_interp[z] + self.radio_abundances_interp[ident]

    def _decay_abundances(self):
        """Determines the decay of all radioactive isotopes. Afterwards their
        mass fractions are added to the stable elements.

        Notes
        -----
        Must be called after _remap_abundances.

        Raises
        ------
        AttributeError if called before _remap_abundances
        """
        try:
            self.radio_abundances_interp
        except AttributeError:
            logger.critical("Abundances must be remapped before decay is"
                            " handled")
            raise AttributeError("no radio_abundances_interp; call "
                                 "_remap_abundances first")

        for i in xrange(self.N_interp):
            comp = {}
            mass = 0
            for ident in self.radio_abundances_interp.keys():
                Xi = self.radio_abundances_interp[ident][i]
                mass += Xi
                comp[nucname.id(ident)] = Xi
            inp = material.Material(comp, mass=mass)
            res = inp.decay(
                (self.t - self.orig.t).to("s").value).mult_by_mass()

            for item in res.items():
                z = nucname.znum(item[0])
                self.abundances_interp[z][i] = \
                    self.abundances_interp[z][i] + item[1]

    def _be_fix(self, to_z=6):
        """Set Beryllium abundance to zero and add its mass fraction to another
        (specified) element. This is done to avoid issue #438.

        Notes
        -----
        Must be called after _remap_abundances.

        Parameters
        ---------
        to_z : int
            proton number of the destination element for the Be fix (default 6,
            i.e. carbon)

        Raises
        ------
        AttributeError if called before _remap_abundances
        """
        try:
            self.abundances_interp
        except AttributeError:
            logger.critical("Abundances must be remapped before the Be fix can"
                            " be applied")
            raise AttributeError("no abundances_interp; call"
                                 " _remap_abundances first")

        logger.info("Total relative Be mass in model: {:e}\n".format(
            (self.abundances_interp[4] * self.dm_interp).sum() /
            self.dm_interp.sum()))

        self.abundances_interp[to_z] = (
            self.abundances_interp[to_z] + self.abundances_interp[4])
        self.abundances_interp[4] = np.zeros(self.N_interp)

    def _write_tardis_abundance_file(self, fname="tardis_abundances.dat"):
        """Write specific density file for Tardis

        Parameters
        ----------
        fname : str
            name of the Tardis specific structure file (default
            'tardis_abundances.dat')
        """
        f = open(fname, "w")
        f.write("# index Z=1 - Z={:d}\n".format(zmax))
        X = np.zeros((zmax+1, self.N_interp+1))

        X[0, :] = np.arange(self.N_interp+1)
        X[1:, 1:] = np.array(
            [self.abundances_interp[z] for z in xrange(1, zmax+1)])
        X[1:, 0] = X[1:, 1]

        np.savetxt(f, X.T, fmt=["% 4d"] + ["%.7e" for i in xrange(1, zmax+1)])
        f.close()

    def _write_tardis_density_file(self, fname="tardis_densities.dat"):
        """Write specific abundance file for Tardis

        Parameters
        ----------
        fname : str
            name of the Tardis specific structure file (default
            'tardis_densities.dat')
        """
        f = open(fname, "w")

        f.write("{:f} {:s}\n".format(self.t.to("day").value, "day"))
        f.write("# index velocity (km/s) density (g/cm^3)\n")
        X = np.array([np.arange(self.N_interp+1),
                     np.insert(self.v_interp_r.to("km/s"), 0,
                               self.v_interp_l.to("km/s")[0]).value,
                     np.insert(self.rho_interp.to("g/cm^3"), 0,
                               self.rho_interp.to("g/cm^3")[0]).value]).T
        np.savetxt(f, X, fmt=["% 4d", "% 9.3f", "%.7e"])

        f.close()


"""
Computes the Heat of Formation at 0 K for a given species
"""

import os
import csv
import numpy as np
from qcelemental import constants as qcc
import autoparse.pattern as app
import autoparse.find as apf
import automol.inchi
import automol.graph
from . import util

# Conversion factors
KJ2KCAL = qcc.conversion_factor('kJ/mol', 'kcal/mol')
EH2KCAL = qcc.conversion_factor('hartree', 'kcal/mol')

# Path  the database files (stored in the thermo src directory)
SRC_PATH = os.path.dirname(os.path.realpath(__file__))

def get_hform_298k_thermp(output_string):
    """
    Obtains deltaHf from thermp output
    """

    # Line pattern containing the DeltaHf value at 298 K
    dhf298_pattern = ('h298 final' +
                      app.one_or_more(app.SPACE) +
                      app.capturing(app.FLOAT))
    dhf298 = float(apf.last_capture(dhf298_pattern, output_string))

    return dhf298

def calc_hform_0k(hzero_mol, hzero_basis, basis, coeff, ref_set):
    """ calculates the heat-of-formation at 0 K
    """

    # Calculate the heat of formation
    dhzero = hzero_mol * EH2KCAL
    for i, spc in enumerate(basis):
        h_basis = get_ref_h(spc, ref_set, 0)
        if h_basis is None:
            h_basis = 0.0
        dhzero += coeff[i] * h_basis * KJ2KCAL
        dhzero -= coeff[i] * hzero_basis[i] * EH2KCAL

    return dhzero

def get_ref_h(species, ref, temp):
    """ gets a reference value
    """

    # Set path and name to thermo database file
    thermodb_name = 'thermodb_{}K.csv'.format(str(int(temp)))
    thermodb_file = os.path.join(SRC_PATH, thermodb_name)

    # Find the energy value for the given species and enery type
    h_species = None
    with open(thermodb_file, 'r') as db_file:
        reader = csv.DictReader(db_file)
        for row in reader:
            if row['inchi'] == species:
                val = row[ref]
                if val == '':
                    val = None
                h_species = float(val)

    return h_species

def select_basis(atom_dct, att=0):
    """
    Given a list of atoms, generates a list of molecules
    that is best suited to serve as a basis for those atoms

    :param atomlist: list of atoms
    :type atomlist: list
    :param att: ???
    :type att: ???

    OUPUT:
    basis    - recommended basis as a list of stoichiometries
    """

    # Determine number of basis species required
    nbasis = len(atom_dct)

    # Get a list of all the atom types in the molecule
    atoms = list(atom_dct.keys())

    # Create list of inchi keys corresponding to basis species
    basis = []
    # H2
    basis.append('InChI=1S/H2/h1H')
    # NH3
    if 'N' in atoms:
        basis.append('InChI=1S/H3N/h1H3')
    # CH4
    if 'C' in atoms:
        basis.append('InChI=1S/CH4/h1H4')
    # H2O
    if 'O' in atoms:
        basis.append('InChI=1S/H2O/h1H2')
    # SO2
    if 'S' in atoms:
        basis.append('InChI=1S/O2S/c1-3-2')
        if not 'O' in atoms:
            basis.append('InChI=1S/H2O/h1H2')

    return basis

def get_reduced_basis(basis_ich, species_formula):
    """
    Form a matrix for a given basis and atomlist
    INPUT:
    input_basis     - ich strings for set of reference molecules
    atomlist  - list of atoms (all atoms that appear
                in basis should be in atomlist)
    OUTPUT:
    mat       - matrix (length of basis by length of atomlist)
                (square if done right)
    """

    # Get the basis formulae list
    basis_formulae = [util.inchi_formula(spc) for spc in basis_ich]

    reduced_basis = []
    for i, basis_formula in enumerate(basis_formulae):
        basis_atom_dict = util.get_atom_counts_dict(basis_formula)
        flag = True
        for key, _ in basis_atom_dict.items():
            if key not in species_formula:
                flag = False

        if flag:
            reduced_basis.append(basis_ich[i])

    return reduced_basis

def calc_coefficients(basis, mol_atom_dict):
    """
    Form a matrix for a given basis and atomlist
    INPUT:
    basis     - basis of molecules
    atomlist  - list of atoms (all atoms that appear
                in basis should be in atomlist)
    OUTPUT:
    mat       - matrix (length of basis by length of atomlist)
                (square if done right)
    """

    # Initialize an natoms x natoms matrix
    nbasis = len(basis)
    basis_mat = np.zeros((nbasis, nbasis))

    # Get the basis formulae list
    basis_formulae = [util.inchi_formula(spc) for spc in basis]
    #basis_atom_dict = [automol.geom.formula(automol.inchi.geom(spc) for spc in basis]
    for spc in basis_formulae:
        basis_atom_dict = util.get_atom_counts_dict(spc)
        for atom in basis_atom_dict:
            if not atom in mol_atom_dict:
                mol_atom_dict[atom] = 0
    # Set the elements of the matrix
    for i, spc in enumerate(basis_formulae):
        basis_atom_dict = util.get_atom_counts_dict(spc)
        basis_vals = []
        for key in mol_atom_dict.keys():
            if key in basis_atom_dict:
                basis_vals.append(basis_atom_dict[key])
            else:
                basis_vals.append(0)
        basis_mat[i] = basis_vals

    #  Transpose
    basis_mat = basis_mat.T

    # Form stoich vector
    stoich_vec = np.zeros(len(mol_atom_dict))
    for i, key in enumerate(mol_atom_dict.keys()):
        stoich_vec[i] = mol_atom_dict[key]

    # Solve C = M^-1 S
    basis_mat = np.linalg.inv(basis_mat)
    coeff = np.dot(basis_mat, stoich_vec)

    return coeff

def stoich(ich):
    """
    Finds the stoichiometry of a molecule
    INPUT:
    ich  -- STR inchii
    OUTPUT:
    stoich -- dictionary with key = STR atomsymbol,
                val = INT number of atomsymbol in molecule
    """

    stoich = {'H': 0}
    gra = automol.inchi.graph(ich)
    atms = automol.graph.atoms(gra)
    for atm in atms:
        stoich['H'] += atms[atm][1]
        if atms[atm][0] in stoich:
            stoich[atms[atm][0]] += 1
        else:
            stoich[atms[atm][0]] = 1
    return stoich

def cbhzed(ich):
    """
    Fragments molecule so that each heavy-atom is a seperate fragment
    INPUT:
    ich --  STR inchii name for molecule
    OUTPUT
    frags -- DIC dictionary with keys as STR inchii name for fragments
    and value as INT their coefficient
    """

    # Graphical info about molecule
    gra = automol.inchi.graph(ich)
    rad_atms = list(automol.graph.sing_res_dom_radical_atom_keys(gra))
    atm_vals = automol.graph.atom_element_valences(gra)
    atms = automol.graph.atoms(gra)

    # Determine CBHzed fragments
    frags = {}
    for atm in atm_vals:
        if atm in rad_atms:
            atm_vals[atm] -= 1
        atm_dic = {0: (atms[atm][0], int(atm_vals[atm]), None)}
        gra = (atm_dic, {})
        frag = automol.graph.inchi(gra)
        _add2dic(frags, frag)

    return _balance_frags(ich, frags)

def cbhone(ich):
    """
    Fragments molecule in a way that conserves each heavy-atom/heavy-atom bond
    INPUT:
    ich --  STR inchii name for molecule
    OUTPUT
    frags -- DIC dictionary with keys as STR inchii name for fragments
    and value as INT their coefficient
    """

    # Graphical info about molecule
    gra = automol.inchi.graph(ich)
    atms = automol.graph.atoms(gra)
    bnd_ords = automol.graph.one_resonance_dominant_bond_orders(gra)
    rad_atms = list(automol.graph.sing_res_dom_radical_atom_keys(gra))
    atm_vals = automol.graph.atom_element_valences(gra)
    adj_atms = automol.graph.atom_neighbor_keys(gra)

    # Determine CBHone fragments
    frags = {}
    for atm in atm_vals:
        for adj in list(adj_atms[atm]):
            if atm > adj:
                vali = atm_vals[atm]
                valj = atm_vals[adj]
                if atm in rad_atms:
                    vali -= 1
                if adj in rad_atms:
                    valj -= 1
                key = frozenset({atm, adj})
                bnd_ord = list(bnd_ords[key])[0]
                vali -= bnd_ord
                valj -= bnd_ord
                atm_dic = {0: (atms[atm][0], int(vali), None),
                           1: (atms[adj][0], int(valj), None)}
                bnd_dic = {frozenset({0, 1}): (1, None)}
                gra = (atm_dic, bnd_dic)
                frag = automol.graph.inchi(gra)
                _add2dic(frags, frag)
    frags = {k: v for k, v in frags.items() if v}

    # Balance
    balance_ = _balance(ich, frags)
    balance_ = {k: v for k, v in balance_.items() if v}

    if balance_:
        newfrags = {}
        zedfrags = cbhzed(ich)
        new = {}
        for frag in frags:
            newfrags[frag] = frags[frag]
            new = cbhzed(frag)
            for n in new:
                _add2dic(newfrags, n, - new[n] * frags[frag])
        if not frags:
            frags = cbhzed(ich)
        for frag in zedfrags:
            if frag in newfrags:
                _add2dic(newfrags, frag, zedfrags[frag])
        frags = newfrags
        frags = {k: v for k, v in frags.items() if v}
        balance_ = _balance(ich, frags)
        balance_ = {k: v for k, v in balance_.items() if v}

    if balance_:
        frags = _balance_frags(ich, frags)

    return frags

def cbhtwo(ich):
    """
    Fragments molecule for each heavy-atom to stay bonded to its adjacent atoms
    INPUT:
    ich --  STR inchii name for molecule
    OUTPUT
    frags -- DIC dictionary with keys as STR inchii name for fragments and
    value as INT their coefficient
    """

    # Graphical info about molecule
    gra = automol.inchi.graph(ich)
    atms = automol.graph.atoms(gra)
    bnd_ords = automol.graph.one_resonance_dominant_bond_orders(gra)
    rad_atms = list(automol.graph.sing_res_dom_radical_atom_keys(gra))
    atm_vals = automol.graph.atom_element_valences(gra)
    adj_atms = automol.graph.atom_neighbor_keys(gra)

    # Determine CBHtwo fragments
    frags = {}
    for atm in atms:
        vali = atm_vals[atm]
        if atm in rad_atms:
            vali -= 1
        # First loop over all atoms of this frag to get saturation of atomi
        for adj in list(adj_atms[atm]):
            key = frozenset({atm, adj})
            bnd_ord = list(bnd_ords[key])[0]
            vali -= bnd_ord
        atm_dic = {0: (atms[atm][0], int(vali), None)}
        bnd_dic = {}
        # Then start adding bonds to the bnddic and atomdic
        j = 0
        for adj in list(adj_atms[atm]):
            j += 1
            valj = atm_vals[adj]
            if adj in rad_atms:
                valj -= 1
            key = frozenset({atm, adj})
            bnd_ord = list(bnd_ords[key])[0]
            valj -= bnd_ord
            atm_dic[j] = (atms[adj][0], int(valj), None)
            bnd_dic[frozenset({0, j})] = (1, None)
        gra = (atm_dic, bnd_dic)
        frag = automol.graph.inchi(gra)
        _add2dic(frags, frag)

    frags = {k: v for k, v in frags.items() if v}

    # Balance
    balance_ = _balance(ich, frags)
    balance_ = {k: v for k, v in balance_.items() if v}
    if balance_:
        newfrags = {}
        onefrags = cbhone(ich)
        new = {}
        for frag in frags:
            newfrags[frag] = frags[frag]
            new = cbhone(frag)
            for n in new:
                _add2dic(newfrags, n, - new[n] * frags[frag])
        if not frags:
            frags = cbhone(ich)
        for frag in onefrags:
            if frag in newfrags:
                _add2dic(newfrags, frag, onefrags[frag])
        frags = newfrags
        frags = {k: v for k, v in frags.items() if v}

        balance_ = _balance(ich, frags)
        balance_ = {k: v for k, v in balance_.items() if v}
        if balance_:
            newfrags = {}
            zedfrags = cbhzed(ich)
            new = {}
            for frag in frags:
                newfrags[frag] = frags[frag]
                new = cbhzed(frag)
                for n in new:
                    _add2dic(newfrags, n, - new[n] * frags[frag])
            if not frags:
                frags = cbhzed(ich)
            for frag in zedfrags:
                if frag in newfrags:
                    _add2dic(newfrags, frag, zedfrags[frag])
            frags = newfrags
            frags = {k: v for k, v in frags.items() if v}

            balance_ = _balance(ich, frags)
            balance_ = {k: v for k, v in balance_.items() if v}
            if balance_:
                frags = _balance_frags(ich, frags)

    return frags

def get_basis(ich):
    formula  = util.inchi_formula(ich)
    atm_dict = util.get_atom_counts_dict(formula)
    return select_basis(atm_dict)

def get_cbhzed(ich):
    return list(cbhzed(ich).keys())

def get_cbhone(ich):
    return list(cbhone(ich).keys())

def get_cbhtwo(ich):
    return list(cbhtwo(ich).keys())

def _add2dic(dic, key, val=1):
    if key in dic:
        dic[key] += val
    else:
        dic[key] = val

def _lhs_rhs(frags):
    rhs = {}
    lhs = {}
    for frag in frags:
        if frags[frag] > 0:
            rhs[frag] = frags[frag]
        elif frags[frag] < 0:
            lhs[frag] = - frags[frag]
    return lhs, rhs

def _print_lhs_rhs(ich, frags):
    lhs, rhs = _lhs_rhs(frags)
    lhsprint = automol.inchi.smiles(ich)
    rhsprint = ''
    for frag in rhs:
        if rhsprint:
            rhsprint += ' +  {:.1f} {} '.format(
                rhs[frag], automol.inchi.smiles(frag))
        else:
            rhsprint = ' {:.1f} {} '.format(
                rhs[frag], automol.inchi.smiles(frag))
    for frag in lhs:
        lhsprint += ' +  {:.1f} {} '.format(
            lhs[frag], automol.inchi.smiles(frag))
    return '{} --> {}'.format(lhsprint, rhsprint)

def _balance(ich, frags):
    stoichs = {}
    for frag in frags:
        _stoich = stoich(frag)
        for atm in _stoich:
            if atm in stoichs:
                stoichs[atm] += _stoich[atm] * frags[frag]
            else:
                stoichs[atm] = _stoich[atm] * frags[frag]
    balance_ = {}
    _stoich = stoich(ich)
    for atom in _stoich:
        if atom in stoichs:
            balance_[atom] = _stoich[atom] - stoichs[atom]
        else:
            balance_[atom] = _stoich[atom]
    balance_ = {x: y for x, y in balance_.items() if y != 0}
    return balance_

def _balance_frags(ich, frags):
    balance_ = _balance(ich, frags)
    methane = automol.smiles.inchi('C')
    water = automol.smiles.inchi('O')
    ammonm = automol.smiles.inchi('N')
    hydrgn = automol.smiles.inchi('[H][H]')
    if 'C' in balance_:
        _add2dic(frags, methane, balance_['C'])
    if 'N' in balance_:
        _add2dic(frags, ammonm, balance_['N'])
    if 'O' in balance_:
        _add2dic(frags, water, balance_['O'])
    balance_ = _balance(ich, frags)
    if 'H' in balance_:
        _add2dic(frags, hydrgn, balance_['H']/2)
    return frags