'''
Original version: Lutz Raestatter Oct 1(?), 2021
Modify to work with flythrough: Oct 5, 2021 (Rebecca Ringuette)
'''
from datetime import datetime,timedelta,timezone

#standard model dictionary for reference
model_varnames={
    ### 3D variables - to be aggregated into 4D ###
    'bx':['B_x','x component of magnetic field',0,'GSE','car',['time','x','y','z'],'nT'],
    'by':['B_y','y component of magnetic field',1,'GSE','car',['time','x','y','z'],'nT'],
    'bz':['B_z','z component of magnetic field',2,'GSE','car',['time','x','y','z'],'nT'],
    'bx1':['B1_x','x component of magnetic field (on grid cell faces)',3,'GSE','car',['time','x','x','x'],'nT'],
    'by1':['B1_y','y component of magnetic field (on grid cell faces)',4,'GSE','car',['time','y','y','y'],'nT'],
    'bz1':['B1_z','z component of magnetic field (on grid cell faces)',5,'GSE','car',['time','z','z','z'],'nT'],
    'ex':['E_x','x component of electric field (on grid cell edges)',6,'GSE','car',['time','x','x','x'],'mV/m'],
    'ey':['E_y','y component of electric field (on grid cell edges)',7,'GSE','car',['time','y','y','y'],'mV/m'],
    'ez':['E_z','z component of electric field (on grid cell edges)',8,'GSE','car',['time','z','z','z'],'mV/m'],
    'vx':['V_x','x component of plasma velocity',9,'GSE','car',['time','x','y','z'],'km/s'],
    'vy':['V_y','y component of plasma velocity',10,'GSE','car',['time','x','y','z'],'km/s'],
    'vz':['V_z','z component of plasma velocity',11,'GSE','car',['time','x','y','z'],'km/s'],
    'rr':['N_plasma','plasma number denstity (hydrogen equivalent)',12,'GSE','car',['time','x','y','z'],'1/cm**3'],
    'resis':['eta','resistivity',13,'GSE','car',['time','x','y','z'],'m**2/s'],
    'pp':['P_plasma','plasma pressure',14,'GSE','car',['time','x','y','z'],'pPa'],
    'xjx':['J_x','x component of current density',15,'GSE','car',['time','x','y','z'],'muA/m**2'],
    'xjy':['J_y','y component of current density',16,'GSE','car',['time','x','y','z'],'muA/m**2'],
    'xjz':['J_z','z component of current density',17,'GSE','car',['time','x','y','z'],'muA/m**2'],
    #add current density components here
}

# variable linkage to grid position vectors are established during variable registration
# these are gx_bx, gy_bx, ... gz_ez affecting magnetic field (b1x,b1y,b1z) and electric field (ex,ey,ez)

#convert an array of timestamps to an array of hrs since midnight
def ts_to_hrs(time_val, filedate):
    '''Convert utc timestamp to hours since midnight on filedate.'''

    return (datetime.utcfromtimestamp(time_val).replace(tzinfo=timezone.utc)-filedate).total_seconds()/3600.

def hrs_to_ts(hrs, filedate):
    '''Add hours to filedate and return utc timestamp.'''

    return datetime.timestamp(filedate+timedelta(hours=float(hrs)))

#sample file name: 'D:/OpenGGCM_GM/Data/Yihua_1/Yihua_Zheng_090721_1.3df_2015-10-16_12.nc'
#main function/class definition
def MODEL():
    from numpy import array, NaN, diff, sqrt, sum, expand_dims, zeros, where
    from time import perf_counter
    import xarray, dask
    from os.path import isfile
    from kamodo import Kamodo, kamodofy
    from kamodo_ccmc.readers.reader_utilities import register_interpolator, define_4d_gridded_interpolator

    class MODEL(Kamodo):
        '''OpenGGCM_GM magnetosphere reader'''
        def __init__(self,full_file_prefix, variables_requested=[], runname = "noname",
                     filetime=False, verbose=False, gridded_int=True, printfiles=False,
                     fulltime=True, missing_value=NaN, **kwargs):
            super(MODEL, self).__init__()
            t0=perf_counter() # profiling time stamp

            #convert files to netcdf4 if needed
            nc_file = full_file_prefix+'.nc'  # input file name: file_dir/YYYY-MM-DD_HH.nc
            if isfile(nc_file):  #file already prepared!
                self.conversion_test = True  #default value
            else:  #file not prepared, prepare it
                try:  #I don't have the file converter, so leave in try/except for now
                    from kamodo_ccmc.readers.openggcm_to_cdf import openggcm_combine_magnetosphere_files as gmconv
                    self.conversion_test = gmconv(full_file_prefix)
                    #should return a boolean (True is successful, False if not)
                except:
                    self.conversion_test = False

            #data are time-wrapped in files. This logic prevents the flythrough from breaking from this decision.
            #cdf_data.added_time_at_beginning and cdf_data.added_time_at_end = 0 if not, 1 if yes
            cdf_data = xarray.open_dataset(nc_file, chunks={'time':100,'x':100,'y':100,'z':100})
            if not fulltime and cdf_data.added_time_at_end:  #use unwrapped time values
                t = cdf_data.variables['_time'].values[:-1]  #skip added time at end
            else: #need wrapped time array for interpolation
                t = cdf_data.variables['_time'].values
            self._time = t

            #establish time attributes first
            self.filedate = datetime.strptime(cdf_data.filedate+' 00:00:00',
                                              '%Y-%m-%d %H:%M:%S').replace(tzinfo=timezone.utc)
            if len(t)>1: self.dt = diff(t).max()*3600.  #t is in hours since midnight
            else: self.dt = 0
            self.datetimes=[datetime.utcfromtimestamp(hrs_to_ts(t[0], self.filedate)).isoformat(sep=' '),
                            datetime.utcfromtimestamp(hrs_to_ts(t[-1], self.filedate)).isoformat(sep=' ')]
            self.filetimes=[hrs_to_ts(t[0], self.filedate), hrs_to_ts(t[-1], self.filedate)]   #timestamps for matching in wrapper

            #return time information only for flythrough
            if filetime:
                return

            #if variables are given as integers, convert to standard names
            if len(variables_requested)>0:
                if isinstance(variables_requested[0], int):
                    print('Integers detected. Converting...', end="")
                    tmp_var = [value[0] for key, value in model_varnames.items()\
                                           if value[2] in variables_requested]
                    variables_requested = tmp_var
                    print('Converted:', variables_requested)

            #perform initial check on variables_requested list
            if len(variables_requested)>0 and fulltime:
                test_list = [value[0] for key, value in model_varnames.items()]
                err_list = [item for item in variables_requested if item not in test_list]
                if len(err_list)>0: print('Variable name(s) not recognized:', err_list)

            #collect variable list
            if len(variables_requested)>0:
                gvar_list = [key for key, value in model_varnames.items() \
                                 if value[0] in variables_requested and \
                                     key in cdf_data.variables.keys()]  # file variable names

                #check for variables requested but not available
                if len(gvar_list)!=len(variables_requested):
                    err_list = [value[0] for key, value in model_varnames.items() \
                                 if value[0] in variables_requested and \
                                     key not in cdf_data.variables.keys()]
                    if len(err_list)>0: print('Some requested variables are not available:', err_list)
            else:  #only input variables on the avoid_list if specifically requested
                avoid_list = []   #empty for now
                gvar_list = [key for key in cdf_data.variables.keys() \
                             if key in model_varnames.keys() and \
                                 key not in avoid_list]

            # Store variable's units and reference to Dataset object in memory
            variables = {model_varnames[key][0]:{'units':model_varnames[key][-1],
                                   'data':getattr(cdf_data, key)}\
                              for key in gvar_list}
            if verbose: print('Done reading in variable data.', full_file_prefix)

            #store variables
            self.near_Earth_boundary_radius = cdf_data.near_Earth_boundary_radius
            self.near_Earth_boundary_radius_unit = cdf_data.near_Earth_boundary_radius_units
            self.missing_value = NaN
            self.verbose = verbose
            self.filename = cdf_data.file.split(',')
            self.modelname = cdf_data.model
            self.runname = runname
            self.modelname = 'OpenGGCM_GM'
            self._registered = 0
            if printfiles:
                print('Files:')
                for file in self.filename: print(file)

            #add coordinate grids as needed
            #grid_list = list of coordinate names in cdf file
            grid_list = ['_x','_y','_z','_x_bx','_y_bx','_z_bx','_x_by','_y_by',
                         '_z_by','_x_bz','_y_bz','_z_bz','_x_ex','_y_ex','_z_ex',
                         '_x_ey','_y_ey','_z_ey','_x_ez','_y_ez','_z_ez']
            #trim down to only save coordinate grids needed for variables requested
            #  and available in file
            if 'B1_x' not in variables.keys():
                grid_list.remove('_x_bx')
                grid_list.remove('_y_bx')
                grid_list.remove('_z_bx')
            if 'B1_y' not in variables.keys():
                grid_list.remove('_x_by')
                grid_list.remove('_y_by')
                grid_list.remove('_z_by')
            if 'B1_z' not in variables.keys():
                grid_list.remove('_x_bz')
                grid_list.remove('_y_bz')
                grid_list.remove('_z_bz')
            if 'E1_x' not in variables.keys():
                grid_list.remove('_x_ex')
                grid_list.remove('_y_ex')
                grid_list.remove('_z_ex')
            if 'E1_y' not in variables.keys():
                grid_list.remove('_x_ey')
                grid_list.remove('_y_ey')
                grid_list.remove('_z_ey')
            if 'E1_z' not in variables.keys():
                grid_list.remove('_x_ez')
                grid_list.remove('_y_ez')
                grid_list.remove('_z_ez')
            for grid in grid_list:
                setattr(self, grid, getattr(cdf_data, grid).values)  #store coordinate data
            cdf_data.close()  #done with file

            #register interpolators for each variable
            varname_list, self.variables = [key for key in variables.keys()], {}  #store original list b/c gridded interpolators
            t_reg = perf_counter()
            for varname in varname_list:  #all are 3D variables
                #make dimension names uniform for easier interpolation later
                dim_names, dims_dict = list(variables[varname]['data'].dims), {}
                dims_dict[dim_names[0]], dims_dict[dim_names[1]] = 'time', 'x'
                dims_dict[dim_names[2]], dims_dict[dim_names[3]] = 'y', 'z'
                variables[varname]['data'] = variables[varname]['data'].rename(
                    **dims_dict)

                #store and register data
                self.variables[varname] = dict(units = variables[varname]['units'],
                                               data = variables[varname]['data'])     #not saving data to decrease memory demand
                self.register_variable(self.variables[varname]['units'],
                                       self.variables[varname]['data'], varname,
                                       gridded_int)
            #cdf_data.close()
            if verbose: print(f'Took {perf_counter()-t_reg:.5f}s to register '+\
                              f'{len(varname_list)} variables.')
            if verbose: print(f'Took a total of {perf_counter()-t0:.5f}s to kamodofy '+\
                              f'{len(gvar_list)} variables.')

        #define and register a 4D variable-----------------------------------------
        def register_variable(self, units, variable, varname, gridded_int):
            x_, y_, z_ = self.get_grid(varname) # variable may have different grid positions in the staggered grid of the model
            xvec_dependencies = {'time':'hr','x':'R_E','y':'R_E','z':'R_E'}

            #variable is a DataArray object with a built-in interpolator, add coordinates
            variable = variable.assign_coords({'time':self._time,'x':x_,'y':y_,'z':z_})

            self = self.custom_interp(units, variable, varname, xvec_dependencies,
                                      gridded_int)   #regdef_4D_interpolators
            return

        def get_grid(self, varname):
            """fetch the grid positon for this variable"""

            if varname == 'B1_x':
                return self._x_bx, self._y_bx, self._z_bx
            elif varname == 'B1_y':
                return self._x_by, self._y_by, self._z_by
            elif varname == 'B1_z':
                return self._x_bz, self._y_bz, self._z_bz
            elif varname == 'E1_x':
                return self._x_ex, self._y_ex, self._z_ex
            elif varname == 'E1_y':
                return self._x_ey, self._y_ey, self._z_ey
            elif varname == 'E1_z':
                return self._x_ez, self._y_ez, self._z_ez
            else: # (default) positions on plasma grid
                return self._x, self._y, self._z

        def custom_interp(self, units, variable, varname,
                          xvec_dependencies, gridded_int):
            '''define interpolator based on xarray's except need inner boundary limit.
            Decrease memory demand by not saving data arrays.'''

            @kamodofy(units=units, data=variable)
            def interpolator(xvec):  #xvec = [[t1,x1,y1,z1],[t2,x2,y2,z2],...]
                """Interpolates 4d variable without a grid"""

                t0 = perf_counter()
                xvec_arr = array(xvec)
                if len(xvec_arr.shape)==1: xvec_arr = expand_dims(xvec_arr, 0)
                r = sqrt(sum(xvec_arr[:,1:]**2, axis=1))  #calculate radius 1D array
                result = zeros(len(r))  #option 2, not much time diff from option 1
                for i in range(len(r)):
                    if r[i]>self.near_Earth_boundary_radius:
                        result[i] = float(variable.interp(time=xvec_arr[i][0], x=xvec_arr[i][1],
                                                          y=xvec_arr[i][2], z=xvec_arr[i][3]).values)
                    else:
                        result[i] = NaN
                print(f'Took {perf_counter()-t0:.3f} s for {len(xvec)} positions.')
                return result

            self = register_interpolator(self, varname, interpolator, xvec_dependencies)

            #define and register the gridded interpolator if desired
            if gridded_int:
                self.variables[varname+'_ijk'] = dict(units = units, data = variable)

                @kamodofy(units=units, data={}, arg_units=xvec_dependencies)
                def gridded_interpolator(time, x, y, z):
                    """Interpolates 4d variable on a grid"""

                    t0 = perf_counter()
                    time, x, y, z = array(time), array(x), array(y), array(z)
                    r = sqrt(x**2+y**2+z**2)
                    return where(r>self.near_Earth_boundary_radius,
                                 variable.interp(time=time, x=x, y=y, z=z).values, NaN)

                self = register_interpolator(self, varname+'_ijk',
                                                         gridded_interpolator,
                                                         xvec_dependencies)
            return

    return MODEL


#!/usr/bin/env python3

"""Core classes and functions for completions code."""

import sys
from typing import Dict
from copy import deepcopy

from functools import reduce

from neutrinomass.tensormethod.core import (
    BOSE,
    FERMI,
    Index,
    IndexedField,
    eps,
    get_dynkin,
)
from neutrinomass.tensormethod.lagrangian import Lagrangian

# from neutrinomass.completions.tikzfeynman import tikz_export

class FieldType(IndexedField):
    """Base class for exotic fields."""

    def __new__(cls, *args, **kwargs):
        return super(FieldType, cls).__new__(cls, *args, **kwargs)

    def lower_su2(self, skip=[]):
        undotted, dotted, _, isospin, _ = self.indices_by_type.values()
        epsilons = []
        partner = self
        for idx in [*undotted, *dotted, *isospin]:
            if idx.index_type in skip:
                continue
            lower = str(idx) + "^"
            partner = partner.substituted_indices((idx, lower))
            epsilon = eps("-" + lower + " -" + str(idx))
            epsilons.append(epsilon)

        # For VectorLikeDiracFermion, keep track of (un)barred boolean
        is_unbarred = None
        if hasattr(self, "is_unbarred"):
            is_unbarred = self.is_unbarred

        return cons_completion_field(partner, is_unbarred=is_unbarred), epsilons

    @property
    def indexed_field(self):
        return IndexedField(label=self.label, indices=self.index_labels)

    # @property
    # def info(self):
    #     return {
    #         "label": self.label,
    #         "indices": self.indices,
    #         "charges": self.charges,
    #         "is_conj": self.is_conj,
    #         "symmetry": self.symmetry,
    #         "latex": self.latex,
    #     }

    def __hash__(self):
        dict_ = {
            "indices": tuple(sorted(self.indices)),
            "label": self.label,
            "dynkin": self.dynkin,
            "charges": tuple(sorted(self.charges.items())),
        }
        return hash(tuple(dict_.items()))

    def __deepcopy__(self, memo):
        return self.__class__(
            label=self.label,
            indices=deepcopy(self.indices, memo),
            charges=deepcopy(self.charges, memo),
            latex=self.latex,
            is_conj=self.is_conj,
            symmetry=deepcopy(self.symmetry, memo),
            comm=self.comm,
        )

class ComplexScalar(FieldType):
    def __init__(
        self,
        label,
        indices,
        charges=None,
        latex=None,
        is_conj=False,
        symmetry=None,
        **kwargs,
    ):
        if isinstance(indices, list):
            indices = " ".join(str(i) for i in indices)

        IndexedField.__init__(
            self,
            label=label,
            indices=indices,
            charges=charges,
            is_conj=is_conj,
            symmetry=None,
            comm=BOSE,
            latex=latex,
            # **kwargs,
        )

        assert self.is_scalar

    @property
    def mass_term(self):
        lower, epsilons = self.conj.lower_su2()
        return reduce(lambda x, y: x * y, epsilons, lower * self)

    # @property
    # def kinetic_term(self):
    #     """I think this will mess with U(1) symmetries algorithm as currently
    #     implemented"""
    #     lower, epsilons = self.conj.lower_su2()
    #     all_og_indices = "u0 d0 " + " ".join(str(i) for i in self.indices)
    #     all_lower_indices = "u1 d1 " + " ".join(str(i) for i in lower.indices)
    #     deriv_og = D(lower, "11")(all_og_indices)
    #     deriv_lower = D(lower, "11")(all_lower_indices)
    #     term = deriv_og * deriv_lower
    #     return reduce(lambda x, y: x * y, epsilons, term)

def assert_real_rep(indices: str, charges) -> None:
    if charges and "y" in charges:
        assert charges["y"] == 0

    colour_dynkin_str = get_dynkin(indices)[2:4]

    # simplistic way to check if colour rep is real
    assert colour_dynkin_str == "".join(reversed(colour_dynkin_str))

class RealScalar(FieldType):
    def __init__(
        self, label, indices, charges=None, latex=None, is_conj=False, **kwargs
    ):
        if isinstance(indices, list):
            indices = " ".join(str(i) for i in indices)

        IndexedField.__init__(
            self,
            label=label,
            indices=indices,
            charges=charges,
            is_conj=is_conj,
            symmetry=None,
            comm=BOSE,
            latex=latex,
            # **kwargs,
        )

        assert self.is_scalar
        assert_real_rep(indices, charges)

    def swap_colour_indices(self):
        """New copy of field with colour indices flipped.

        # TODO Refactor this out with majorana_partner to FieldType
        """
        undotted, dotted, colour, isospin, _ = Index.indices_by_type(
            self.indices
        ).values()
        colour = tuple(i.conj for i in colour)
        indices = undotted + dotted + colour + isospin

        return RealScalar(
            self.label,
            indices,
            latex=self.latex,
            is_conj=self.is_conj,
            symmetry=self.symmetry,
            comm=BOSE,
            charges=None,
        )

    @property
    def mass_term(self):
        lower, epsilons = self.swap_colour_indices().lower_su2()
        return reduce(lambda x, y: x * y, epsilons, lower * self)

class MajoranaFermion(FieldType):
    def __init__(
        self, label, indices, charges=None, latex=None, is_conj=False, **kwargs
    ):
        if isinstance(indices, list):
            indices = " ".join(str(i) for i in indices)

        IndexedField.__init__(
            self,
            label=label,
            indices=indices,
            charges=charges,
            is_conj=is_conj,
            symmetry=None,
            comm=FERMI,
            latex=latex,
            # **kwargs,
        )

        assert self.is_fermion
        assert_real_rep(indices, charges)

    def majorana_partner(self):
        """New copy of fermion with colour indices flipped"""
        undotted, dotted, colour, isospin, _ = Index.indices_by_type(
            self.indices
        ).values()
        colour = tuple(i.conj for i in colour)
        indices = undotted + dotted + colour + isospin

        return MajoranaFermion(
            self.label,
            indices,
            latex=self.latex,
            is_conj=self.is_conj,
            symmetry=self.symmetry,
            comm=FERMI,
            charges=None,
        )

    @property
    def conj_indices(self):
        return self.conj

    @property
    def mass_term(self):
        lower, epsilons = self.majorana_partner().lower_su2()
        return reduce(lambda x, y: x * y, epsilons, lower * self)

class VectorLikeDiracFermion(FieldType):
    """Stands in for two fermion fields distinguished here only by dotted and
    undotted indices.

    Example:
        >>> psi = VectorLikeDiracFermion("ψ", "u0 i1")
        >>> psi
        ψ(u0, i1)
        >>> psi.dirac_partner()
        ψ~(u0, i1)

    Here ψ(u0, i1) and ψ~(u0, i1) are different fields.

    """

    def __init__(
        self,
        label,
        indices,
        charges=None,
        latex=None,
        is_unbarred=True,
        is_conj=False,
        symmetry=None,
        comm=FERMI,
        **kwargs,
    ):
        if isinstance(indices, list):
            indices = " ".join(str(i) for i in indices)

        IndexedField.__init__(
            self,
            label=label,
            indices=indices,
            charges=charges,
            is_conj=is_conj,
            symmetry=None,
            latex=latex,
            comm=FERMI,
            # **kwargs,
        )

        assert self.is_fermion
        self.is_unbarred = is_unbarred

    @property
    def conj(self):
        """Returns a copy of self but conjugated"""

        is_conj = self.is_conj
        if is_conj:
            label = self.label.replace("†", "")
        else:
            label = self.label + "†"

        return self.__class__(
            label=label,
            indices=" ".join(i.conj.label for i in self.indices),
            charges={k: -v for k, v in self.charges.items()},
            is_conj=(not is_conj),
            is_unbarred=self.is_unbarred,
            symmetry=self.symmetry,
            comm=self.comm,
        )

    @property
    def conj_indices(self):
        """Returns a copy of self conjugated but leaves charges alone."""
        is_conj = self.is_conj
        if is_conj:
            label = self.label.replace("†", "")
        else:
            label = self.label + "†"

        return self.__class__(
            label=label,
            indices=" ".join(i.conj.label for i in self.indices),
            charges=self.charges,
            is_conj=(not is_conj),
            is_unbarred=self.is_unbarred,
            symmetry=self.symmetry,
            comm=self.comm,
        )

    def dirac_partner(self):
        undotted, dotted, colour, isospin, _ = Index.indices_by_type(
            self.indices
        ).values()
        colour = tuple(i.conj for i in colour)
        indices = undotted + dotted + colour + isospin
        charges = {k: -v for k, v in self.charges.items()}

        is_unbarred = self.is_unbarred
        symb = "~"  # indicate bar with circumflex
        if is_unbarred:
            label = self.label + symb
        else:
            label = self.label.replace(symb, "")

        return VectorLikeDiracFermion(
            label,
            indices,
            charges=charges,
            latex=self.latex,
            is_unbarred=(not self.is_unbarred),
            is_conj=self.is_conj,
            symmetry=self.symmetry,
            comm=FERMI,
        )

    @property
    def mass_term(self):
        lower, epsilons = self.dirac_partner().lower_su2()
        return reduce(lambda x, y: x * y, epsilons, lower * self)

class EffectiveOperator:
    def __init__(self, name, operator):
        self.name = name
        self.operator = operator

    @property
    def fields(self):
        return self.operator.fields

    @property
    def indexed_fields(self):
        return self.operator.indexed_fields

    @property
    def mass_dimension(self):
        d = sum(f.mass_dim for f in self.fields)
        d += sum(f.derivs for f in self.fields)
        d_int = int(d)
        assert d == d_int
        return d_int

    @property
    def topology_type(self):
        """Returns a dictionary {"n_scalars": n_scalars, "n_fermions": n_fermions}."""

        n_scalars, n_fermions = 0, 0
        for f in self.fields:
            if f.is_boson:
                n_scalars += 1
            elif f.is_fermion:
                n_fermions += 1

        return {"n_scalars": n_scalars, "n_fermions": n_fermions}

    def __hash__(self):
        return hash((self.name, self.operator.simplify()))

class FailedCompletion:
    def __init__(self, reason: str = ""):
        self.reason = reason

class Completion:
    def __init__(self, operator, partition, graph, exotics, terms, topology=None):
        self.operator = operator
        self.partition = partition
        self.graph = graph
        self.exotics = exotics
        self.terms = terms
        self.topology = topology

    def __eq__(self, other):
        if not isinstance(other, Completion):
            return False
        return (
            self.operator == other.operator
            and self.exotic_info() == other.exotic_info()
            and self.partition == other.partition
        )

    def __hash__(self):
        return hash((self.operator, self.exotic_info(), self.partition))

    def __deepcopy__(self, memo):
        return self.__class__(
            operator=deepcopy(self.operator, memo),
            partition=deepcopy(self.partition, memo),
            graph=deepcopy(self.graph, memo),
            exotics=deepcopy(self.exotics, memo),
            terms=deepcopy(self.terms, memo),
        )

    @property
    def lagrangian(self):
        return Lagrangian(exotics=self.exotics, interaction_terms=self.terms)

    def draw_diagram(self):
        import matplotlib.pyplot as plt
        from matplotlib import rc
        import networkx as nx

        # Unicode Greek edge labels won't render in TeX
        rc("text", usetex=False)

        g = self.graph
        edge_labels = nx.get_edge_attributes(g, name="particle")
        pos = nx.spring_layout(g)

        plt.figure()
        nx.draw(g, pos=pos, edge_color="black", node_size=0)
        nx.draw_networkx_edge_labels(
            g, pos=pos, edge_labels=edge_labels, font_color="red"
        )
        plt.axis("off")
        plt.show()

    def exotic_info(self) -> Dict[FieldType, tuple]:
        info = {}
        for e in self.exotics:
            # normalise hypercharge to be positive
            if e.y < 0:
                charges = sorted(e.conj.charges.items())
                sm = e.conj.sm_irrep
            else:
                charges = sorted(e.charges.items())
                sm = e.sm_irrep

            if e.is_fermion:
                lorentz = "F"
            elif e.is_scalar:
                lorentz = "S"
            else:
                raise ValueError("Unrecognised exotic field type.")

            info[e] = (lorentz,) + sm + tuple(charges)

        return info

    def exotic_fields(self):
        return set([e.field for e in self.exotics])

    def info(self):
        print("Fields:")
        for k, v in self.exotic_info().items():
            print("{:<5s}{:<20s}".format(k.label, format_quantum_numbers(v)))

        print("\nLagrangian:")
        in_jupyter = sys.argv[-1].endswith("json")
        for term in self.terms:
            if in_jupyter:
                display(term)
            else:
                print(term)

        print("\nDiagram:")
        self.draw_diagram()

class Model:
    def __init__(self, completions):
        self.completions = completions

    @property
    def exotic_numbers(self):
        return sorted(
            set(
                format_quantum_numbers(i)
                for i in self.completions[0].exotic_info().values()
            )
        )

    def __repr__(self):
        return "Model(" + " + ".join(i for i in self.exotic_numbers) + ")"

def format_quantum_numbers(info: tuple):
    """Takes an expression like

    ('S', 1, 0, 2, ('3b', 1), ('y', 2/3))

    and returns a string like

    S(3, 3, 2/3)(3b: 1)

    """
    lorentz, su3_up, su3_down, su2, *charges = info
    su3_dim = lambda m, n: 0.5 * (m + 1) * (n + 1) * (m + n + 2)
    # For now just add the bar for more lowered than raised indices, but for
    # larger reps this will be problematic
    su3_dim_format = lambda m, n: str(int(su3_dim(m, n))) + ("b" if n > m else "")
    charges_dict = dict(charges)
    return f"{lorentz}({su3_dim_format(int(su3_up), int(su3_down))}, {str(int(su2) + 1)}, {charges_dict['y']})({charges_dict['3b']})"

def cons_completion_field(indexed_field: IndexedField, is_unbarred=None) -> FieldType:
    label = indexed_field.label
    indices = indexed_field.indices
    charges = indexed_field.charges
    latex = indexed_field.latex
    is_conj = indexed_field.is_conj

    if indexed_field.is_fermion:
        if indexed_field.is_real_sm_irrep:
            return MajoranaFermion(
                label, indices, charges=charges, latex=latex, is_conj=is_conj
            )

        if is_unbarred is None:
            is_unbarred = True

        return VectorLikeDiracFermion(
            label,
            indices,
            charges=charges,
            latex=latex,
            is_unbarred=is_unbarred,
            is_conj=is_conj,
        )

    if indexed_field.is_scalar:
        if indexed_field.is_real_sm_irrep:
            return RealScalar(label, indices, charges=charges, latex=latex)

        return ComplexScalar(
            label, indices, charges=charges, latex=latex, is_conj=is_conj
        )

    raise Exception("Unrecognised Lorentz structure in field.")


"""
BPASS version 2

Eldridge, Stanway et al, 2017, PASA 34, 58

Paper describes version 2.1. Latest updates in 2.2 described in
Stanway & Eldridge (2018).
"""

import re, os
import numpy as np
from scipy.interpolate import interp1d
from ares.physics.Constants import h_p, c, erg_per_ev, g_per_msun, s_per_yr, \
    s_per_myr, m_H, Lsun

_input = os.getenv('ARES') + '/input/bpass_v2/SEDS'

metallicities = \
{
 '040': 0.040, '030': 0.040, '020': 0.020, '010': 0.010,
 '008': 0.008, '006': 0.006, '004': 0.004, '003': 0.003,
 '001': 0.002, '001': 0.001,
}

info = \
{
 'flux_units': r'$L_{\odot} \ \AA^{-1}$',
}

_log10_times = np.arange(6, 10.1, 0.1)
times = 10**_log10_times / 1e6            # Convert from yr to Myr

def _kwargs_to_fn(**kwargs):
    """
    Determine filename of appropriate BPASS lookup table based on kwargs.
    """

    # All files share this prefix
    fn = 'spectra'

    assert kwargs['source_ssp'], \
        "No support for continuous star formation in BPASS v2."
    assert kwargs['source_nebular'] in [0, 2], \
        "No support for nebular emission in BPASS v2."

    if kwargs['source_binaries']:
        fn += '-bin'
    else:
        fn += '-sin'

    fn += '-imf{}'.format(str((kwargs['source_imf'] - 1)).replace('.', ''))
    fn += '_{}'.format(str(int(kwargs['source_imf_Mmax'])))

    # Metallicity
    fn += '.z{!s}'.format(str(int(kwargs['source_Z'] * 1e3)).zfill(3))

    if kwargs['source_sed_degrade'] is not None:
        fn += '.deg{}'.format(kwargs['source_sed_degrade'])

    fn += '.dat'

    return _input + '/' + fn

def _load(**kwargs):
    """
    Return wavelengths, fluxes, for given set of parameters (at all times).
    """

    Zvals_l = list(metallicities.values())
    Zvals = np.sort(Zvals_l)

    # Interpolate
    if kwargs['source_Z'] not in Zvals_l:
        tmp = kwargs.copy()

        _fn = []
        spectra = []
        del tmp['source_Z']
        for Z in Zvals:
            _w1, _d1, fn = _load(source_Z=Z, **tmp)
            spectra.append(_d1.copy())
            _fn.append(fn)

        wavelengths = wave = _w1
        data = spectra

    # No interpolation necessary
    else:
        fn = _fn = _kwargs_to_fn(**kwargs)

        _raw_data = np.loadtxt(fn)

        data = np.array(_raw_data[:,1:])
        wavelengths = _raw_data[:,0]

        data *= Lsun

    return wavelengths, data, _fn


#!/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.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#         Timothy Berkelbach <tim.berkelbach@gmail.com>
#

'''
QCISD for real integrals
8-fold permutation symmetry has been used
(ij|kl) = (ji|kl) = (kl|ij) = ...
'''

import numpy
from pyscf import gto
from pyscf import lib
from pyscf.lib import logger
from pyscf.cc.ccsd import CCSD, _add_vvvv, _flops, _ChemistsERIs
from pyscf import __config__

BLKMIN = getattr(__config__, 'cc_ccsd_blkmin', 4)
MEMORYMIN = getattr(__config__, 'cc_ccsd_memorymin', 2000)

# t1: ia
# t2: ijab
def kernel(mycc, eris=None, t1=None, t2=None, max_cycle=50, tol=1e-8,
           tolnormt=1e-6, verbose=None):
    log = logger.new_logger(mycc, verbose)
    if eris is None:
        eris = mycc.ao2mo(mycc.mo_coeff)
    if t1 is None and t2 is None:
        t1, t2 = mycc.get_init_guess(eris)
    elif t2 is None:
        t2 = mycc.get_init_guess(eris)[1]

    cput1 = cput0 = (logger.process_clock(), logger.perf_counter())
    eold = 0
    eccsd = mycc.energy(t1, t2, eris)
    log.info('Init E_corr(QCISD) = %.15g', eccsd)

    if isinstance(mycc.diis, lib.diis.DIIS):
        adiis = mycc.diis
    elif mycc.diis:
        adiis = lib.diis.DIIS(mycc, mycc.diis_file, incore=mycc.incore_complete)
        adiis.space = mycc.diis_space
    else:
        adiis = None

    conv = False
    for istep in range(max_cycle):
        t1new, t2new = mycc.update_amps(t1, t2, eris)
        tmpvec = mycc.amplitudes_to_vector(t1new, t2new)
        tmpvec -= mycc.amplitudes_to_vector(t1, t2)
        normt = numpy.linalg.norm(tmpvec)
        tmpvec = None
        if mycc.iterative_damping < 1.0:
            alpha = mycc.iterative_damping
            t1new = (1-alpha) * t1 + alpha * t1new
            t2new *= alpha
            t2new += (1-alpha) * t2
        t1, t2 = t1new, t2new
        t1new = t2new = None
        t1, t2 = mycc.run_diis(t1, t2, istep, normt, eccsd-eold, adiis)
        eold, eccsd = eccsd, mycc.energy(t1, t2, eris)
        log.info('cycle = %d  E_corr(QCISD) = %.15g  dE = %.9g  norm(t1,t2) = %.6g',
                 istep+1, eccsd, eccsd - eold, normt)
        cput1 = log.timer('QCISD iter', *cput1)
        if abs(eccsd-eold) < tol and normt < tolnormt:
            conv = True
            break
    log.timer('QCISD', *cput0)
    return conv, eccsd, t1, t2

def update_amps(mycc, t1, t2, eris):
    assert(isinstance(eris, _ChemistsERIs))

    time0 = logger.process_clock(), logger.perf_counter()
    log = logger.Logger(mycc.stdout, mycc.verbose)
    nocc, nvir = t1.shape
    fock = eris.fock
    mo_e_o = eris.mo_energy[:nocc]
    mo_e_v = eris.mo_energy[nocc:] + mycc.level_shift

    t1new = numpy.zeros_like(t1)
    t2new = _add_vvvv(mycc, 0*t1, t2, eris, t2sym='jiba')
    t2new *= .5  # *.5 because t2+t2.transpose(1,0,3,2) in the end
    time1 = log.timer_debug1('vvvv', *time0)

#** make_inter_F
    fov = fock[:nocc,nocc:].copy()
    t1new += fov

    foo = fock[:nocc,:nocc] - numpy.diag(mo_e_o)
    fvv = fock[nocc:,nocc:] - numpy.diag(mo_e_v)

    if mycc.incore_complete:
        fswap = None
    else:
        fswap = lib.H5TmpFile()
    fwVOov, fwVooV = _add_ovvv_(mycc, t1, t2, eris, fvv, t1new, t2new, fswap)
    time1 = log.timer_debug1('ovvv', *time1)

    woooo = numpy.asarray(eris.oooo).transpose(0,2,1,3).copy()

    unit = nocc**2*nvir*7 + nocc**3 + nocc*nvir**2
    mem_now = lib.current_memory()[0]
    max_memory = max(0, mycc.max_memory - mem_now)
    blksize = min(nvir, max(BLKMIN, int((max_memory*.9e6/8-nocc**4)/unit)))
    log.debug1('max_memory %d MB,  nocc,nvir = %d,%d  blksize = %d',
               max_memory, nocc, nvir, blksize)

    for p0, p1 in lib.prange(0, nvir, blksize):
        wVOov = fwVOov[p0:p1]
        wVooV = fwVooV[p0:p1]
        eris_ovoo = eris.ovoo[:,p0:p1]
        eris_oovv = numpy.empty((nocc,nocc,p1-p0,nvir))
        def load_oovv(p0, p1):
            eris_oovv[:] = eris.oovv[:,:,p0:p1]
        with lib.call_in_background(load_oovv, sync=not mycc.async_io) as prefetch_oovv:
            #:eris_oovv = eris.oovv[:,:,p0:p1]
            prefetch_oovv(p0, p1)
            wVOov -= lib.einsum('jbik,ka->bjia', eris_ovoo, t1)
            t2new[:,:,p0:p1] += wVOov.transpose(1,2,0,3)
            eris_ovoo = None
        load_oovv = prefetch_oovv = None

        wVOov *= 0 # QCI

        eris_ovvo = numpy.empty((nocc,p1-p0,nvir,nocc))
        def load_ovvo(p0, p1):
            eris_ovvo[:] = eris.ovvo[:,p0:p1]
        with lib.call_in_background(load_ovvo, sync=not mycc.async_io) as prefetch_ovvo:
            #:eris_ovvo = eris.ovvo[:,p0:p1]
            prefetch_ovvo(p0, p1)
            t1new[:,p0:p1] -= numpy.einsum('jb,jiab->ia', t1, eris_oovv)
            wVooV -= eris_oovv.transpose(2,0,1,3)
            wVOov += wVooV*.5  #: bjia + bija*.5
        load_ovvo = prefetch_ovvo = None

        t2new[:,:,p0:p1] += (eris_ovvo*0.5).transpose(0,3,1,2)
        eris_voov = eris_ovvo.conj().transpose(1,0,3,2)
        t1new[:,p0:p1] += 2*numpy.einsum('jb,aijb->ia', t1, eris_voov)
        eris_ovvo = None
        eris_oovv = tmp = None

        fov[:,p0:p1] += numpy.einsum('kc,aikc->ia', t1, eris_voov) * 2
        fov[:,p0:p1] -= numpy.einsum('kc,akic->ia', t1, eris_voov)

        tau = t2[:,:,p0:p1]
        theta  = tau.transpose(1,0,2,3) * 2
        theta -= tau
        fvv -= lib.einsum('cjia,cjib->ab', theta.transpose(2,1,0,3), eris_voov)
        foo += lib.einsum('aikb,kjab->ij', eris_voov, theta)
        theta = None
        woooo += lib.einsum('ijab,aklb->ijkl', tau, eris_voov)
        tau = None

        def update_wVooV(q0, q1, tau):
            wVooV[:] += lib.einsum('bkic,jkca->bija', eris_voov[:,:,:,q0:q1], tau)
        with lib.call_in_background(update_wVooV, sync=not mycc.async_io) as update_wVooV:
            for q0, q1 in lib.prange(0, nvir, blksize):
                tau  = t2[:,:,q0:q1] * .5
                #:wVooV += lib.einsum('bkic,jkca->bija', eris_voov[:,:,:,q0:q1], tau)
                update_wVooV(q0, q1, tau)
        tau = update_wVooV = None
        def update_t2(q0, q1, tmp):
            t2new[:,:,q0:q1] += tmp.transpose(2,0,1,3)
            tmp *= .5
            t2new[:,:,q0:q1] += tmp.transpose(0,2,1,3)
        with lib.call_in_background(update_t2, sync=not mycc.async_io) as update_t2:
            for q0, q1 in lib.prange(0, nvir, blksize):
                tmp = lib.einsum('jkca,ckib->jaib', t2[:,:,p0:p1,q0:q1], wVooV)
                #:t2new[:,:,q0:q1] += tmp.transpose(2,0,1,3)
                #:tmp *= .5
                #:t2new[:,:,q0:q1] += tmp.transpose(0,2,1,3)
                update_t2(q0, q1, tmp)
                tmp = None

        wVOov += eris_voov
        eris_VOov = -.5 * eris_voov.transpose(0,2,1,3)
        eris_VOov += eris_voov
        eris_voov = None
        def update_wVOov(q0, q1, tau):
            wVOov[:,:,:,q0:q1] += .5 * lib.einsum('aikc,kcjb->aijb', eris_VOov, tau)
        with lib.call_in_background(update_wVOov, sync=not mycc.async_io) as update_wVOov:
            for q0, q1 in lib.prange(0, nvir, blksize):
                tau  = t2[:,:,q0:q1].transpose(1,3,0,2) * 2
                tau -= t2[:,:,q0:q1].transpose(0,3,1,2)
                #:wVOov[:,:,:,q0:q1] += .5 * lib.einsum('aikc,kcjb->aijb', eris_VOov, tau)
                update_wVOov(q0, q1, tau)
                tau = None
        def update_t2(q0, q1, theta):
            t2new[:,:,q0:q1] += lib.einsum('kica,ckjb->ijab', theta, wVOov)
        with lib.call_in_background(update_t2, sync=not mycc.async_io) as update_t2:
            for q0, q1 in lib.prange(0, nvir, blksize):
                theta  = t2[:,:,p0:p1,q0:q1] * 2
                theta -= t2[:,:,p0:p1,q0:q1].transpose(1,0,2,3)
                #:t2new[:,:,q0:q1] += lib.einsum('kica,ckjb->ijab', theta, wVOov)
                update_t2(q0, q1, theta)
                theta = None
        eris_VOov = wVOov = wVooV = update_wVOov = None
        time1 = log.timer_debug1('voov [%d:%d]'%(p0, p1), *time1)
    fwVOov = fwVooV = fswap = None

    for p0, p1 in lib.prange(0, nvir, blksize):
        theta = t2[:,:,p0:p1].transpose(1,0,2,3) * 2 - t2[:,:,p0:p1]
        t1new += numpy.einsum('jb,ijba->ia', fov[:,p0:p1], theta)
        t1new -= lib.einsum('jbki,kjba->ia', eris.ovoo[:,p0:p1], theta)

        tau = t2[:,:,p0:p1]
        t2new[:,:,p0:p1] += .5 * lib.einsum('ijkl,klab->ijab', woooo, tau)
        theta = tau = None

    t2new += lib.einsum('ijac,bc->ijab', t2, fvv)
    t2new -= lib.einsum('ki,kjab->ijab', foo, t2)

    eia = mo_e_o[:,None] - mo_e_v
    t1new += numpy.einsum('ib,ab->ia', t1, fvv)
    t1new -= numpy.einsum('ja,ji->ia', t1, foo)
    t1new /= eia

    #: t2new = t2new + t2new.transpose(1,0,3,2)
    for i in range(nocc):
        if i > 0:
            t2new[i,:i] += t2new[:i,i].transpose(0,2,1)
            t2new[i,:i] /= lib.direct_sum('a,jb->jab', eia[i], eia[:i])
            t2new[:i,i] = t2new[i,:i].transpose(0,2,1)
        t2new[i,i] = t2new[i,i] + t2new[i,i].T
        t2new[i,i] /= lib.direct_sum('a,b->ab', eia[i], eia[i])

    time0 = log.timer_debug1('update t1 t2', *time0)
    return t1new, t2new

def _add_ovvv_(mycc, t1, t2, eris, fvv, t1new, t2new, fswap):
    time1 = logger.process_clock(), logger.perf_counter()
    log = logger.Logger(mycc.stdout, mycc.verbose)
    nocc, nvir = t1.shape
    nvir_pair = nvir * (nvir+1) // 2

    if fswap is None:
        wVOov = numpy.zeros((nvir,nocc,nocc,nvir))
    else:
        wVOov = fswap.create_dataset('wVOov', (nvir,nocc,nocc,nvir), 'f8')
    wooVV = numpy.zeros((nocc,nocc*nvir_pair))

    max_memory = mycc.max_memory - lib.current_memory()[0]
    unit = nocc*nvir**2*3 + nocc**2*nvir + 2
    blksize = min(nvir, max(BLKMIN, int((max_memory*.95e6/8-wooVV.size)/unit)))
    if not mycc.direct:
        unit = nocc*nvir**2*3 + nocc**2*nvir + 2 + nocc*nvir**2 + nocc*nvir
        blksize = min(nvir, max(BLKMIN, int((max_memory*.95e6/8-wooVV.size-nocc**2*nvir)/unit)))
    log.debug1('max_memory %d MB,  nocc,nvir = %d,%d  blksize = %d',
               max_memory, nocc, nvir, blksize)

    def load_ovvv(buf, p0):
        if p0 < nvir:
            p1 = min(nvir, p0+blksize)
            buf[:p1-p0] = eris.ovvv[:,p0:p1].transpose(1,0,2)

    with lib.call_in_background(load_ovvv, sync=not mycc.async_io) as prefetch:
        buf = numpy.empty((blksize,nocc,nvir_pair))
        buf_prefetch = numpy.empty((blksize,nocc,nvir_pair))

        load_ovvv(buf_prefetch, 0)
        for p0, p1 in lib.prange(0, nvir, blksize):
            buf, buf_prefetch = buf_prefetch, buf
            prefetch(buf_prefetch, p1)

            eris_vovv = buf[:p1-p0]
            eris_vovv = lib.unpack_tril(eris_vovv.reshape((p1-p0)*nocc,nvir_pair))
            eris_vovv = eris_vovv.reshape(p1-p0,nocc,nvir,nvir)

            wVOov[p0:p1] = lib.einsum('biac,jc->bija', eris_vovv, t1)

            theta = t2[:,:,p0:p1].transpose(1,2,0,3) * 2
            theta -= t2[:,:,p0:p1].transpose(0,2,1,3)
            t1new += lib.einsum('icjb,cjba->ia', theta, eris_vovv)
            theta = None
            time1 = log.timer_debug1('vovv [%d:%d]'%(p0, p1), *time1)

    if fswap is None:
        wooVV = lib.unpack_tril(wooVV.reshape(nocc**2,nvir_pair))
        return wVOov, wooVV.reshape(nocc,nocc,nvir,nvir).transpose(2,1,0,3)
    else:
        fswap.create_dataset('wVooV', (nvir,nocc,nocc,nvir), 'f8')
        wooVV = wooVV.reshape(nocc,nocc,nvir_pair)
        tril2sq = lib.square_mat_in_trilu_indices(nvir)
        for p0, p1 in lib.prange(0, nvir, blksize):
            fswap['wVooV'][p0:p1] = wooVV[:,:,tril2sq[p0:p1]].transpose(2,1,0,3)
        return fswap['wVOov'], fswap['wVooV']

def energy(mycc, t1=None, t2=None, eris=None):
    '''CCSD correlation energy'''
    if t1 is None: t1 = mycc.t1
    if t2 is None: t2 = mycc.t2
    if eris is None: eris = mycc.ao2mo()

    nocc, nvir = t1.shape
    fock = eris.fock
    e = numpy.einsum('ia,ia', fock[:nocc,nocc:], t1) * 2
    max_memory = mycc.max_memory - lib.current_memory()[0]
    blksize = int(min(nvir, max(BLKMIN, max_memory*.3e6/8/(nocc**2*nvir+1))))
    for p0, p1 in lib.prange(0, nvir, blksize):
        eris_ovvo = eris.ovvo[:,p0:p1]
        tau = t2[:,:,p0:p1]
        e += 2 * numpy.einsum('ijab,iabj', tau, eris_ovvo)
        e -=     numpy.einsum('jiab,iabj', tau, eris_ovvo)
    if abs(e.imag) > 1e-4:
        logger.warn(mycc, 'Non-zero imaginary part found in QCISD energy %s', e)
    return e.real

def as_scanner(cc):
    '''Generating a scanner/solver for QCISD PES.

    The returned solver is a function. This function requires one argument
    "mol" as input and returns total QCISD energy.

    '''
    if isinstance(cc, lib.SinglePointScanner):
        return cc

    logger.info(cc, 'Set %s as a scanner', cc.__class__)

    class QCISD_Scanner(cc.__class__, lib.SinglePointScanner):
        def __init__(self, cc):
            self.__dict__.update(cc.__dict__)
            self._scf = cc._scf.as_scanner()
        def __call__(self, mol_or_geom, **kwargs):
            if isinstance(mol_or_geom, gto.Mole):
                mol = mol_or_geom
            else:
                mol = self.mol.set_geom_(mol_or_geom, inplace=False)

            if self.t2 is not None:
                last_size = self.vector_size()
            else:
                last_size = 0

            self.reset(mol)

            mf_scanner = self._scf
            mf_scanner(mol)
            self.mo_coeff = mf_scanner.mo_coeff
            self.mo_occ = mf_scanner.mo_occ
            if last_size != self.vector_size():
                self.t1 = self.t2 = None
            self.kernel(self.t1, self.t2, **kwargs)
            return self.e_tot
    return QCISD_Scanner(cc)

class QCISD(CCSD):
    '''restricted QCISD
    '''

    def dump_flags(self, verbose=None):
        log = logger.new_logger(self, verbose)
        log.info('')
        log.info('******** %s ********', self.__class__)
        log.info('QCISD nocc = %s, nmo = %s', self.nocc, self.nmo)
        if self.frozen is not None:
            log.info('frozen orbitals %s', self.frozen)
        log.info('max_cycle = %d', self.max_cycle)
        log.info('direct = %d', self.direct)
        log.info('conv_tol = %g', self.conv_tol)
        log.info('conv_tol_normt = %s', self.conv_tol_normt)
        log.info('diis_space = %d', self.diis_space)
        #log.info('diis_file = %s', self.diis_file)
        log.info('diis_start_cycle = %d', self.diis_start_cycle)
        log.info('diis_start_energy_diff = %g', self.diis_start_energy_diff)
        log.info('max_memory %d MB (current use %d MB)',
                 self.max_memory, lib.current_memory()[0])
        if (log.verbose >= logger.DEBUG1 and
            self.__class__ == QCISD):
            nocc = self.nocc
            nvir = self.nmo - self.nocc
            flops = _flops(nocc, nvir)
            log.debug1('total FLOPs %s', flops)
        return self

    energy = energy
    _add_vvvv = _add_vvvv
    update_amps = update_amps

    def kernel(self, t1=None, t2=None, eris=None):
        return self.qcisd(t1, t2, eris)
    def qcisd(self, t1=None, t2=None, eris=None):
        assert(self.mo_coeff is not None)
        assert(self.mo_occ is not None)

        if self.verbose >= logger.WARN:
            self.check_sanity()
        self.dump_flags()

        if eris is None:
            eris = self.ao2mo(self.mo_coeff)

        self.e_hf = getattr(eris, 'e_hf', None)
        if self.e_hf is None:
            self.e_hf = self._scf.e_tot

        self.converged, self.e_corr, self.t1, self.t2 = \
                kernel(self, eris, t1, t2, max_cycle=self.max_cycle,
                       tol=self.conv_tol, tolnormt=self.conv_tol_normt,
                       verbose=self.verbose)
        self._finalize()
        return self.e_corr, self.t1, self.t2

    as_scanner = as_scanner

    def qcisd_t(self, t1=None, t2=None, eris=None):
        from pyscf.cc import qcisd_t
        if t1 is None: t1 = self.t1
        if t2 is None: t2 = self.t2
        if eris is None: eris = self.ao2mo(self.mo_coeff)
        return qcisd_t.kernel(self, eris, t1, t2, self.verbose)

RQCISD = QCISD

if __name__ == '__main__':
    from pyscf import scf

    mol = gto.Mole()
    #mol.atom = [
    #    [8 , (0. , 0.     , 0.)],
    #    [1 , (0. , -0.757 , 0.587)],
    #    [1 , (0. , 0.757  , 0.587)]]
    mol.atom = [['Ne', (0,0,0)]]
    mol.basis = 'cc-pvdz'
    mol.verbose = 7
    mol.spin = 0
    mol.build()
    mf = scf.RHF(mol).run(conv_tol=1e-14)

    mycc = QCISD(mf, frozen=1)
    ecc, t1, t2 = mycc.kernel()
    et = mycc.qcisd_t()
    print("QCISD(T) =", mycc.e_tot+et)

    mol = gto.Mole()
    mol.atom = """C  0.000  0.000  0.000
                  H  0.637  0.637  0.637
                  H -0.637 -0.637  0.637
                  H -0.637  0.637 -0.637
                  H  0.637 -0.637 -0.637"""
    mol.basis = 'cc-pvdz'
    mol.verbose = 7
    mol.spin = 0
    mol.build()
    mf = scf.RHF(mol).run(conv_tol=1e-14)

    mycc = QCISD(mf, frozen=1)
    ecc, t1, t2 = mycc.kernel()
    print(mycc.e_tot - -40.383989)
    et = mycc.qcisd_t()
    print(mycc.e_tot+et - -40.387679)


#!/usr/bin/env python

"""
m2g.track
~~~~~~~~~~

Contains m2g's fiber reconstruction and tractography functionality.
Theory described here: https://neurodata.io/talks/ndmg.pdf#page=21
"""

# system imports
import os

# external package imports
import numpy as np
import nibabel as nib

# dipy imports
from dipy.tracking.streamline import Streamlines
from dipy.tracking import utils
from dipy.tracking.local_tracking import LocalTracking
from dipy.tracking.local_tracking import ParticleFilteringTracking
from dipy.tracking.stopping_criterion import BinaryStoppingCriterion
from dipy.tracking.stopping_criterion import ActStoppingCriterion
from dipy.tracking.stopping_criterion import CmcStoppingCriterion

from dipy.reconst.dti import fractional_anisotropy, TensorModel, quantize_evecs
from dipy.reconst.shm import CsaOdfModel
from dipy.reconst.csdeconv import ConstrainedSphericalDeconvModel, recursive_response

from dipy.data import get_sphere
from dipy.direction import peaks_from_model, ProbabilisticDirectionGetter
from m2g.utils.gen_utils import timer

from m2g.stats import qa_tensor

def build_seed_list(mask_img_file, stream_affine, dens):
    """uses dipy tractography utilities in order to create a seed list for tractography

    Parameters
    ----------
    mask_img_file : str
        path to mask of area to generate seeds for
    stream_affine : ndarray
        4x4 array with 1s diagonally and 0s everywhere else
    dens : int
        seed density

    Returns
    -------
    ndarray
        locations for the seeds
    """

    mask_img = nib.load(mask_img_file)
    mask_img_data = mask_img.get_data().astype("bool")
    seeds = utils.random_seeds_from_mask(
        mask_img_data,
        affine=stream_affine,
        seeds_count=int(dens),
        seed_count_per_voxel=True,
    )
    return seeds

def tens_mod_fa_est(gtab, dwi_file, B0_mask):
    """Estimate a tensor FA image to use for registrations using dipy functions

    Parameters
    ----------
    gtab : GradientTable
        gradient table created from bval and bvec file
    dwi_file : str
        Path to eddy-corrected and RAS reoriented dwi image
    B0_mask : str
        Path to nodif B0 mask (averaged b0 mask)

    Returns
    -------
    str
        Path to tensor_fa image file
    """

    data = nib.load(dwi_file).get_fdata()

    print("Generating simple tensor FA image to use for registrations...")
    nodif_B0_img = nib.load(B0_mask)
    B0_mask_data = nodif_B0_img.get_fdata().astype("bool")
    nodif_B0_affine = nodif_B0_img.affine
    model = TensorModel(gtab)
    mod = model.fit(data, B0_mask_data)
    FA = fractional_anisotropy(mod.evals)
    FA[np.isnan(FA)] = 0
    fa_img = nib.Nifti1Image(FA.astype(np.float32), nodif_B0_affine)
    fa_path = f"{os.path.dirname(B0_mask)}/tensor_fa.nii.gz"
    nib.save(fa_img, fa_path)
    return fa_path

class RunTrack:
    def __init__(
        self,
        dwi_in,
        nodif_B0_mask,
        gm_in_dwi,
        vent_csf_in_dwi,
        csf_in_dwi,
        wm_in_dwi,
        gtab,
        mod_type,
        track_type,
        mod_func,
        qa_tensor_out,
        seeds,
        stream_affine,
    ):
        """A class for deterministic tractography in native space

        Parameters
        ----------
        dwi_in : str
            path to the input dwi image to perform tractography on.
            Should be a nifti, gzipped nifti, or other image that nibabel
            is capable of reading, with data as a 4D object.
        nodif_B0_mask : str
            path to the mask of the b0 mean volume. Should be a nifti,
            gzipped nifti, or other image file that nibabel is capable of
            reading, with data as a 3D object.
        gm_in_dwi : str
            Path to gray matter segmentation in EPI space. Should be a nifti,
            gzipped nifti, or other image file that nibabel is capable of
            reading, with data as a 3D object
        vent_csf_in_dwi : str
            Ventricular CSF Mask in EPI space. Should be a nifti,
            gzipped nifti, or other image file that nibabel is capable of
            reading, with data as a 3D object
        csf_in_dwi : str
            Path to CSF mask in EPI space. Should be a nifti, gzipped nifti, or other image file that nibabel compatable
        wm_in_dwi : str
            Path to white matter probabilities in EPI space. Should be a nifti,
            gzipped nifti, or other image file that nibabel is capable of
            reading, with data as a 3D object.
        gtab : gradient table
            gradient table created from bval and bvec files
        mod_type : str
            Determinstic (det) or probabilistic (prob) tracking
        track_type : str
            Tracking approach: local or particle
        mod_func : str
            Diffusion model: csd or csa
        qa_tensor: str
            path to store the qa for tensor/directions of model
        seeds : ndarray
            ndarray of seeds for tractography
        stream_affine : ndarray
            4x4 2D array with 1s diagonaly and 0s everywhere else
        """

        self.dwi = dwi_in
        self.nodif_B0_mask = nodif_B0_mask
        self.gm_in_dwi = gm_in_dwi
        self.vent_csf_in_dwi = vent_csf_in_dwi
        self.csf_in_dwi = csf_in_dwi
        self.wm_in_dwi = wm_in_dwi
        self.gtab = gtab
        self.mod_type = mod_type
        self.track_type = track_type
        self.qa_tensor_out = qa_tensor_out
        self.seeds = seeds
        self.mod_func = mod_func
        self.stream_affine = stream_affine

    @timer
    def run(self):
        """Creates the tracktography tracks using dipy commands and the specified tracking type and approach

        Returns
        -------
        ArraySequence
            contains the tractography track raw data for further analysis

        Raises
        ------
        ValueError
            Raised when no seeds are supplied or no valid seeds were found in white-matter interface
        ValueError
            Raised when no seeds are supplied or no valid seeds were found in white-matter interface
        """
        self.tiss_classifier = self.prep_tracking()
        if self.mod_type == "det":
            if self.mod_func == "csa":
                self.mod = self.odf_mod_est()
            elif self.mod_func == "csd":
                self.mod = self.csd_mod_est()
            if self.track_type == "local":
                tracks = self.local_tracking()
            elif self.track_type == "particle":
                tracks = self.particle_tracking()
            else:
                raise ValueError(
                    "Error: Either no seeds supplied, or no valid seeds found in white-matter interface"
                )
        elif self.mod_type == "prob":
            if self.mod_func == "csa":
                self.mod = self.odf_mod_est()
            elif self.mod_func == "csd":
                self.mod = self.csd_mod_est()
            if self.track_type == "local":
                tracks = self.local_tracking()
            elif self.track_type == "particle":
                tracks = self.particle_tracking()
        else:
            raise ValueError(
                "Error: Either no seeds supplied, or no valid seeds found in white-matter interface"
            )
        tracks = Streamlines([track for track in tracks if len(track) > 60])
        return tracks

    @staticmethod
    def make_hdr(streamlines, hdr):
        trk_hdr = nib.streamlines.trk.TrkFile.create_empty_header()
        trk_hdr["hdr_size"] = 1000
        trk_hdr["dimensions"] = hdr["dim"][1:4].astype("float32")
        trk_hdr["voxel_sizes"] = hdr["pixdim"][1:4]
        trk_hdr["voxel_to_rasmm"] = np.eye(4)
        trk_hdr["voxel_order"] = "RAS"
        trk_hdr["pad2"] = "RAS"
        trk_hdr["image_orientation_patient"] = np.array(
            [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
        ).astype("float32")
        trk_hdr["endianness"] = "<"
        trk_hdr["_offset_data"] = 1000
        trk_hdr["nb_streamlines"] = streamlines.total_nb_rows

        return trk_hdr

    def prep_tracking(self):
        """Uses nibabel and dipy functions in order to load the grey matter, white matter, and csf masks
        and use a tissue classifier (act, cmc, or binary) on the include/exclude maps to make a tissueclassifier object

        Returns
        -------
        ActStoppingCriterion, CmcStoppingCriterion, or BinaryStoppingCriterion
            The resulting tissue classifier object, depending on which method you use (currently only does act)
        """

        if self.track_type == "local":
            tiss_class = "bin"
        elif self.track_type == "particle":
            tiss_class = "cmc"

        self.dwi_img = nib.load(self.dwi)
        self.data = self.dwi_img.get_data()
        # Loads mask and ensures it's a true binary mask
        self.mask_img = nib.load(self.nodif_B0_mask)
        self.mask = self.mask_img.get_data() > 0
        # Load tissue maps and prepare tissue classifier
        self.gm_mask = nib.load(self.gm_in_dwi)
        self.gm_mask_data = self.gm_mask.get_data()
        self.wm_mask = nib.load(self.wm_in_dwi)
        self.wm_mask_data = self.wm_mask.get_data()
        self.wm_in_dwi_data = nib.load(self.wm_in_dwi).get_data().astype("bool")
        if tiss_class == "act":
            self.vent_csf_in_dwi = nib.load(self.vent_csf_in_dwi)
            self.vent_csf_in_dwi_data = self.vent_csf_in_dwi.get_data()
            self.background = np.ones(self.gm_mask.shape)
            self.background[
                (self.gm_mask_data + self.wm_mask_data + self.vent_csf_in_dwi_data) > 0
            ] = 0
            self.include_map = self.wm_mask_data
            self.include_map[self.background > 0] = 0
            self.exclude_map = self.vent_csf_in_dwi_data
            self.tiss_classifier = ActStoppingCriterion(
                self.include_map, self.exclude_map
            )
        elif tiss_class == "bin":
            self.tiss_classifier = BinaryStoppingCriterion(self.wm_in_dwi_data)
            # self.tiss_classifier = BinaryStoppingCriterion(self.mask)
        elif tiss_class == "cmc":
            self.vent_csf_in_dwi = nib.load(self.vent_csf_in_dwi)
            self.vent_csf_in_dwi_data = self.vent_csf_in_dwi.get_data()
            voxel_size = np.average(self.wm_mask.get_header()["pixdim"][1:4])
            step_size = 0.2
            self.tiss_classifier = CmcStoppingCriterion.from_pve(
                self.wm_mask_data,
                self.gm_mask_data,
                self.vent_csf_in_dwi_data,
                step_size=step_size,
                average_voxel_size=voxel_size,
            )
        else:
            pass
        return self.tiss_classifier

    @timer
    def tens_mod_est(self):

        print("Fitting tensor model...")
        self.model = TensorModel(self.gtab)
        self.ten = self.model.fit(self.data, self.wm_in_dwi_data)
        self.fa = self.ten.fa
        self.fa[np.isnan(self.fa)] = 0
        self.sphere = get_sphere("repulsion724")
        self.ind = quantize_evecs(self.ten.evecs, self.sphere.vertices)
        return self.ten

    @timer
    def odf_mod_est(self):

        print("Fitting CSA ODF model...")
        self.mod = CsaOdfModel(self.gtab, sh_order=6)
        return self.mod

    @timer
    def csd_mod_est(self):

        print("Fitting CSD model...")
        try:
            print("Attempting to use spherical harmonic basis first...")
            self.mod = ConstrainedSphericalDeconvModel(self.gtab, None, sh_order=6)
        except:
            print("Falling back to estimating recursive response...")
            self.response = recursive_response(
                self.gtab,
                self.data,
                mask=self.wm_in_dwi_data,
                sh_order=6,
                peak_thr=0.01,
                init_fa=0.08,
                init_trace=0.0021,
                iter=8,
                convergence=0.001,
                parallel=False,
            )
            print("CSD Reponse: " + str(self.response))
            self.mod = ConstrainedSphericalDeconvModel(self.gtab, self.response,sh_order=6)
        return self.mod

    @timer
    def local_tracking(self):

        self.sphere = get_sphere("repulsion724")
        if self.mod_type == "det":
            print("Obtaining peaks from model...")
            self.mod_peaks = peaks_from_model(
                self.mod,
                self.data,
                self.sphere,
                relative_peak_threshold=0.5,
                min_separation_angle=25,
                mask=self.wm_in_dwi_data,
                npeaks=5,
                normalize_peaks=True,
            )
            qa_tensor.create_qa_figure(self.mod_peaks.peak_dirs, self.mod_peaks.peak_values, self.qa_tensor_out, self.mod_func)
            self.streamline_generator = LocalTracking(
                self.mod_peaks,
                self.tiss_classifier,
                self.seeds,
                self.stream_affine,
                step_size=0.5,
                return_all=True,
            )
        elif self.mod_type == "prob":
            print("Preparing probabilistic tracking...")
            print("Fitting model to data...")
            self.mod_fit = self.mod.fit(self.data, self.wm_in_dwi_data)
            print("Building direction-getter...")
            self.mod_peaks = peaks_from_model(
                self.mod,
                self.data,
                self.sphere,
                relative_peak_threshold=0.5,
                min_separation_angle=25,
                mask=self.wm_in_dwi_data,
                npeaks=5,
                normalize_peaks=True,
            )
            qa_tensor.create_qa_figure(self.mod_peaks.peak_dirs, self.mod_peaks.peak_values, self.qa_tensor_out, self.mod_func)
            try:
                print(
                    "Proceeding using spherical harmonic coefficient from model estimation..."
                )
                self.pdg = ProbabilisticDirectionGetter.from_shcoeff(
                    self.mod_fit.shm_coeff, max_angle=60.0, sphere=self.sphere
                )
            except:
                print("Proceeding using FOD PMF from model estimation...")
                self.fod = self.mod_fit.odf(self.sphere)
                self.pmf = self.fod.clip(min=0)
                self.pdg = ProbabilisticDirectionGetter.from_pmf(
                    self.pmf, max_angle=60.0, sphere=self.sphere
                )
            self.streamline_generator = LocalTracking(
                self.pdg,
                self.tiss_classifier,
                self.seeds,
                self.stream_affine,
                step_size=0.5,
                return_all=True,
            )
        print("Reconstructing tractogram streamlines...")
        self.streamlines = Streamlines(self.streamline_generator)
        return self.streamlines

    @timer
    def particle_tracking(self):

        self.sphere = get_sphere("repulsion724")
        if self.mod_type == "det":
            maxcrossing = 1
            print("Obtaining peaks from model...")
            self.mod_peaks = peaks_from_model(
                self.mod,
                self.data,
                self.sphere,
                relative_peak_threshold=0.5,
                min_separation_angle=25,
                mask=self.wm_in_dwi_data,
                npeaks=5,
                normalize_peaks=True,
            )
            qa_tensor.create_qa_figure(self.mod_peaks.peak_dirs, self.mod_peaks.peak_values, self.qa_tensor_out, self.mod_func)
            self.streamline_generator = ParticleFilteringTracking(
                self.mod_peaks,
                self.tiss_classifier,
                self.seeds,
                self.stream_affine,
                max_cross=maxcrossing,
                step_size=0.5,
                maxlen=1000,
                pft_back_tracking_dist=2,
                pft_front_tracking_dist=1,
                particle_count=15,
                return_all=True,
            )
        elif self.mod_type == "prob":
            maxcrossing = 2
            print("Preparing probabilistic tracking...")
            print("Fitting model to data...")
            self.mod_fit = self.mod.fit(self.data, self.wm_in_dwi_data)
            print("Building direction-getter...")
            self.mod_peaks = peaks_from_model(
                self.mod,
                self.data,
                self.sphere,
                relative_peak_threshold=0.5,
                min_separation_angle=25,
                mask=self.wm_in_dwi_data,
                npeaks=5,
                normalize_peaks=True,
            )
            qa_tensor.create_qa_figure(self.mod_peaks.peak_dirs, self.mod_peaks.peak_values, self.qa_tensor_out, self.mod_func)
            try:
                print(
                    "Proceeding using spherical harmonic coefficient from model estimation..."
                )
                self.pdg = ProbabilisticDirectionGetter.from_shcoeff(
                    self.mod_fit.shm_coeff, max_angle=60.0, sphere=self.sphere
                )
            except:
                print("Proceeding using FOD PMF from model estimation...")
                self.fod = self.mod_fit.odf(self.sphere)
                self.pmf = self.fod.clip(min=0)
                self.pdg = ProbabilisticDirectionGetter.from_pmf(
                    self.pmf, max_angle=60.0, sphere=self.sphere
                )
            self.streamline_generator = ParticleFilteringTracking(
                self.pdg,
                self.tiss_classifier,
                self.seeds,
                self.stream_affine,
                max_cross=maxcrossing,
                step_size=0.5,
                maxlen=1000,
                pft_back_tracking_dist=2,
                pft_front_tracking_dist=1,
                particle_count=15,
                return_all=True,
            )
        print("Reconstructing tractogram streamlines...")
        self.streamlines = Streamlines(self.streamline_generator)
        return self.streamlines


from typing import List

import numpy

import torch

def jaccard(_box_a, _box_b):
    # 计算真实框的左上角和右下角
    b1_x1, b1_x2 = _box_a[:, 0] - _box_a[:, 2] / 2, _box_a[:, 0] + _box_a[:, 2] / 2
    b1_y1, b1_y2 = _box_a[:, 1] - _box_a[:, 3] / 2, _box_a[:, 1] + _box_a[:, 3] / 2
    # 计算先验框的左上角和右下角
    b2_x1, b2_x2 = _box_b[:, 0] - _box_b[:, 2] / 2, _box_b[:, 0] + _box_b[:, 2] / 2
    b2_y1, b2_y2 = _box_b[:, 1] - _box_b[:, 3] / 2, _box_b[:, 1] + _box_b[:, 3] / 2
    box_a = torch.zeros_like(_box_a)
    box_b = torch.zeros_like(_box_b)
    box_a[:, 0], box_a[:, 1], box_a[:, 2], box_a[:, 3] = b1_x1, b1_y1, b1_x2, b1_y2
    box_b[:, 0], box_b[:, 1], box_b[:, 2], box_b[:, 3] = b2_x1, b2_y1, b2_x2, b2_y2
    A = box_a.size(0)
    B = box_b.size(0)
    max_xy = torch.min(box_a[:, 2:].unsqueeze(1).expand(A, B, 2),
                       box_b[:, 2:].unsqueeze(0).expand(A, B, 2))
    min_xy = torch.max(box_a[:, :2].unsqueeze(1).expand(A, B, 2),
                       box_b[:, :2].unsqueeze(0).expand(A, B, 2))
    inter = torch.clamp((max_xy - min_xy), min=0)

    inter = inter[:, :, 0] * inter[:, :, 1]
    # 计算先验框和真实框各自的面积
    area_a = ((box_a[:, 2] - box_a[:, 0]) *
              (box_a[:, 3] - box_a[:, 1])).unsqueeze(1).expand_as(inter)  # [A,B]
    area_b = ((box_b[:, 2] - box_b[:, 0]) *
              (box_b[:, 3] - box_b[:, 1])).unsqueeze(0).expand_as(inter)  # [A,B]
    # 求IOU
    union = area_a + area_b - inter
    return inter / union  # [A,B]

def intersect(box_a, box_b):
    """ We resize both tensors to [A,B,2] without new malloc:
    [A,2] -> [A,1,2] -> [A,B,2]
    [B,2] -> [1,B,2] -> [A,B,2]
    Then we compute the area of intersect between box_a and box_b.
    Args:
      box_a: (tensor) bounding boxes, Shape: [A,4].
      box_b: (tensor) bounding boxes, Shape: [B,4].
    Return:
      (tensor) intersection area, Shape: [A,B].
    """
    A = box_a.size(0)
    B = box_b.size(0)
    max_xy = torch.min(box_a[:, 2:].unsqueeze(1).expand(A, B, 2),
                       box_b[:, 2:].unsqueeze(0).expand(A, B, 2))
    min_xy = torch.max(box_a[:, :2].unsqueeze(1).expand(A, B, 2),
                       box_b[:, :2].unsqueeze(0).expand(A, B, 2))
    inter = torch.clamp((max_xy - min_xy), min=0)
    return inter[:, :, 0] * inter[:, :, 1]
    # inter[:, :, 0] is the width of intersection and inter[:, :, 1] is height

def jaccard_tensor(box_a: torch.Tensor, box_b: torch.Tensor) -> torch.Tensor:
    """Compute the jaccard overlap of two sets of boxes.  The jaccard overlap
    is simply the intersection over union of two boxes.  Here we operate on
    ground truth boxes and default boxes.
    E.g.:
        A ∩ B / A ∪ B = A ∩ B / (area(A) + area(B) - A ∩ B)
    Args:
        box_a: (tensor) Ground truth bounding boxes, Shape: [A,4]
        box_b: (tensor) Prior boxes from priorbox layers, Shape: [B,4]
    Return:
        jaccard overlap: (tensor) Shape: [A, B]
    """
    inter = intersect(box_a, box_b)
    area_a = ((box_a[:, 2] - box_a[:, 0]) *
              (box_a[:, 3] - box_a[:, 1])).unsqueeze(1).expand_as(inter)  # [A,B]
    area_b = ((box_b[:, 2] - box_b[:, 0]) *
              (box_b[:, 3] - box_b[:, 1])).unsqueeze(0).expand_as(inter)  # [A,B]
    union = area_a + area_b - inter
    return inter / union  # [A,B]

class YoloV3Loss(torch.nn.Module):
    """
    YoloV3 损失函数
    """

    def __init__(self, config: dict) -> None:
        super().__init__()

        self.lambda_xy = 0.05  # 预测框中心误差权重
        self.lambda_wh = 0.05  # 预测框大小误差权重
        self.lambda_noobj = 1.0  # 预测框置信度误差权重
        self.lambda_obj = 1.0  # 预测框置信度误差权重
        self.lambda_class = 0.5  # 预测框类别误差权重
        self.lambda_conf = 1.0  # 预测框类别误差权重

        self.normd_anchors = numpy.asarray(config["anchors"]).astype(numpy.float32)
        self.normd_anchors[:, :, 0] /= config["image_width"]
        self.normd_anchors[:, :, 1] /= config["image_height"]
        self.normd_anchors = self.normd_anchors.reshape((9, 2))
        self.normd_anchors_box = torch.cat(
            (
                torch.zeros((self.normd_anchors.shape[0], 2)),
                torch.from_numpy(self.normd_anchors)
            ), 1)

        self.classes = config["classes"]
        self.bbox_attrs = 4 + 1 + self.classes

        self.ignore_threshold = 0.3  # iou 忽略的阈值

        self.cuda = config["cuda"]

    def decode_pyramid_boxes(self,
                             pyramid_boxes_list: List[torch.Tensor],
                             pyramid_features: int,
                             predict_feature: torch.Tensor,
                             ignore_on: bool
                             ) -> (
            torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor,
            torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor,
            torch.Tensor, torch.Tensor
    ):
        """
        将真值框分成三个特征层，并变换成 tensor 的格式
        """
        assert pyramid_features in [13, 26, 52]

        if pyramid_features == 13:
            pyramid_anch_index_list = [0, 1, 2]
            cur_anchors = self.normd_anchors[0:3]
        elif pyramid_features == 26:
            pyramid_anch_index_list = [3, 4, 5]
            cur_anchors = self.normd_anchors[3:6]
        elif pyramid_features == 52:
            pyramid_anch_index_list = [6, 7, 8]
            cur_anchors = self.normd_anchors[6:9]
        else:
            raise Exception("unexpected error")

        batch_size = len(pyramid_boxes_list)

        # 将真值框变换为如下的 tensor 格式
        boxes_x = torch.zeros(batch_size, 3, pyramid_features, pyramid_features)
        boxes_y = torch.zeros(batch_size, 3, pyramid_features, pyramid_features)
        boxes_w = torch.zeros(batch_size, 3, pyramid_features, pyramid_features)
        boxes_h = torch.zeros(batch_size, 3, pyramid_features, pyramid_features)

        boxes_loss_weight_xw = torch.zeros(batch_size, 3, pyramid_features, pyramid_features)
        boxes_loss_weight_yh = torch.zeros(batch_size, 3, pyramid_features, pyramid_features)

        boxes_obj_conf = torch.zeros(batch_size, 3, pyramid_features, pyramid_features)
        boxes_class_conf_list = torch.zeros(batch_size, 3, pyramid_features, pyramid_features,
                                            self.classes)

        boxes_obj_mask = torch.zeros(batch_size, 3, pyramid_features, pyramid_features)
        boxes_noobj_mask = torch.ones(batch_size, 3, pyramid_features, pyramid_features)

        # 遍历这一个批次所有的图片
        for bs_i, pyramid_boxes in enumerate(pyramid_boxes_list):
            if pyramid_boxes.shape[0] == 0:
                continue

            # 将真值框的 xy 变换为相对于网格的偏移量，顺便获取真值框在 tensor 表示中的网格索引（truth_grid_x，truth_grid_y——
            truth_feature_box = pyramid_boxes[:, 0:4] * pyramid_features

            truth_grid_x = torch.floor(truth_feature_box[:, 0]).int()
            truth_grid_y = torch.floor(truth_feature_box[:, 1]).int()

            truth_x = truth_feature_box[:, 0] - truth_grid_x
            truth_y = truth_feature_box[:, 1] - truth_grid_y

            # 和真值框 iou 最大的 anchor 索引，确定真值框所在的特征层
            truth_box = pyramid_boxes[:, :4].clone().detach()
            truth_box[:, 0] = 0
            truth_box[:, 1] = 0
            normd_anch_ious = jaccard_tensor(truth_box, self.normd_anchors_box)
            max_anch_ious_index = torch.argmax(normd_anch_ious, dim=-1)

            for box_i, anch_i in enumerate(max_anch_ious_index):
                if anch_i not in pyramid_anch_index_list:
                    continue
                pyramid_anch_i = anch_i % 3

                boxes_x[bs_i, pyramid_anch_i, truth_grid_y[box_i], truth_grid_x[box_i]] = truth_x[box_i]
                boxes_y[bs_i, pyramid_anch_i, truth_grid_y[box_i], truth_grid_x[box_i]] = truth_y[box_i]

                truth_w_box = torch.log(pyramid_boxes[box_i][2] / self.normd_anchors[anch_i][0])
                truth_h_box = torch.log(pyramid_boxes[box_i][3] / self.normd_anchors[anch_i][1])
                boxes_w[bs_i, pyramid_anch_i, truth_grid_y[box_i], truth_grid_x[box_i]] = truth_w_box
                boxes_h[bs_i, pyramid_anch_i, truth_grid_y[box_i], truth_grid_x[box_i]] = truth_h_box

                boxes_loss_weight_xw[bs_i, pyramid_anch_i, truth_grid_y[box_i], truth_grid_x[box_i]] = \
                    pyramid_boxes[box_i][2]
                boxes_loss_weight_yh[bs_i, pyramid_anch_i, truth_grid_y[box_i], truth_grid_x[box_i]] = \
                    pyramid_boxes[box_i][3]

                boxes_obj_conf[bs_i, pyramid_anch_i, truth_grid_y[box_i], truth_grid_x[box_i]] = 1
                boxes_class_conf_list[bs_i, pyramid_anch_i, truth_grid_y[box_i], truth_grid_x[box_i],
                                      pyramid_boxes[box_i][4].int()] = 1

                boxes_obj_mask[bs_i, pyramid_anch_i, truth_grid_y[box_i], truth_grid_x[box_i]] = 1
                boxes_noobj_mask[bs_i, pyramid_anch_i, truth_grid_y[box_i], truth_grid_x[box_i]] = 0  # 可以确定一定有物体

        # print("loss in cuda") if self.cuda else print("loss not in cuda")

        if ignore_on:
            # ----------------------------------------------------------------------------------------------------- #
            # 一些预测框和真值框 iou 较大的地方，有可能有物体
            cur_anchors_num = 3
            predict_feature_height = pyramid_features
            predict_feature_width = pyramid_features
            # 4. 将预测网络输出的特征层进行维度变换，将预测框个数与预测属性分开，并将预测属性转置为末位维度的属性，便于提取和解析
            predict_feature = predict_feature.contiguous().view(
                batch_size,
                cur_anchors_num,
                self.bbox_attrs,
                predict_feature_height,
                predict_feature_width,
            ).permute(0, 1, 3, 4, 2).contiguous()

            # 5. 分隔预测属性
            predict_x = predict_feature[..., 0]
            predict_y = predict_feature[..., 1]
            predict_w = predict_feature[..., 2]
            predict_h = predict_feature[..., 3]

            # 6. 解析 xy
            norm_predict_x = torch.sigmoid(predict_x)
            norm_predict_y = torch.sigmoid(predict_y)
            # 6.1 构造 grid tensor
            grid_x = torch.linspace(0, predict_feature_width - 1, predict_feature_width) \
                .repeat(predict_feature_height, 1) \
                .repeat(batch_size * cur_anchors_num, 1, 1) \
                .view(predict_x.shape)
            grid_y = torch.linspace(0, predict_feature_height - 1, predict_feature_height) \
                .repeat(predict_feature_width, 1) \
                .t() \
                .repeat(batch_size * cur_anchors_num, 1, 1) \
                .view(predict_y.shape)
            if self.cuda:
                grid_x = grid_x.cuda()
                grid_y = grid_y.cuda()
            # 6.2 叠加 grid tensor
            grid_predict_x = norm_predict_x + grid_x
            grid_predict_y = norm_predict_y + grid_y
            # 6.3 归一化 x，y
            normd_predict_x = grid_predict_x / predict_feature_width
            normd_predict_y = grid_predict_y / predict_feature_height

            # 7. 解析 wh
            # 7.1 构造 anchor tensor
            anchor_width = torch.Tensor(cur_anchors)[:, 0].unsqueeze(dim=1)
            anchor_height = torch.Tensor(cur_anchors)[:, 1].unsqueeze(dim=1)
            grid_anchor_width = anchor_width.repeat(batch_size, 1). \
                repeat(1, 1, predict_feature_height * predict_feature_width). \
                view(predict_w.shape)
            grid_anchor_height = anchor_height.repeat(batch_size, 1). \
                repeat(1, 1, predict_feature_height * predict_feature_width). \
                view(predict_h.shape)
            if self.cuda:
                grid_anchor_width = grid_anchor_width.cuda()
                grid_anchor_height = grid_anchor_height.cuda()
            # 7.2 乘以 anchor tensor
            anchord_predict_width = torch.exp(predict_w) * grid_anchor_width
            anchord_predict_height = torch.exp(predict_h) * grid_anchor_height
            # 6.3 归一化 w, h
            normd_predict_w = anchord_predict_width / predict_feature_width
            normd_predict_h = anchord_predict_height / predict_feature_height

            normd_predict_boxes = torch.cat(
                [
                    normd_predict_x.unsqueeze(dim=4),
                    normd_predict_y.unsqueeze(dim=4),
                    normd_predict_w.unsqueeze(dim=4),
                    normd_predict_h.unsqueeze(dim=4),
                ], dim=-1
            )

            for bs_i, pyramid_boxes in enumerate(pyramid_boxes_list):
                if self.cuda:
                    pyramid_boxes = pyramid_boxes.cuda()
                bs_normd_predict_boxes = normd_predict_boxes[bs_i].view(-1, 4)
                predict_truth_ious = jaccard(pyramid_boxes[..., :4], bs_normd_predict_boxes)
                predict_truth_ious_max, _ = torch.max(predict_truth_ious, dim=0)
                predict_truth_ious_max = predict_truth_ious_max.view(normd_predict_boxes[bs_i].size()[:3])
                # a = predict_truth_ious_max > self.ignore_threshold
                # aa = torch.unique(a)
                # print(aa.size())
                boxes_noobj_mask[bs_i][predict_truth_ious_max > self.ignore_threshold] = 0

        if self.cuda:
            return boxes_x.cuda(), \
                   boxes_y.cuda(), \
                   boxes_w.cuda(), \
                   boxes_h.cuda(), \
                   boxes_loss_weight_xw.cuda(), \
                   boxes_loss_weight_yh.cuda(), \
                   boxes_obj_conf.cuda(), \
                   boxes_class_conf_list.cuda(), \
                   boxes_obj_mask.cuda(), \
                   boxes_noobj_mask.cuda()

        return boxes_x, \
               boxes_y, \
               boxes_w, \
               boxes_h, \
               boxes_loss_weight_xw, \
               boxes_loss_weight_yh, \
               boxes_obj_conf, \
               boxes_class_conf_list, \
               boxes_obj_mask, \
               boxes_noobj_mask

    def compute_loss(self, predict_feature: torch.Tensor, decoded_boxes) -> (
            torch.Tensor, torch.Tensor):
        """
        逐个特征层计算损失
        """
        (boxes_x, boxes_y, boxes_w, boxes_h, boxes_loss_weight_xw, boxes_loss_weight_yh, boxes_obj_conf,
         boxes_class_conf_list, boxes_obj_mask, boxes_noobj_mask) = decoded_boxes

        predict_feature = predict_feature.view(
            predict_feature.shape[0],
            3,
            self.bbox_attrs,
            predict_feature.shape[2],
            predict_feature.shape[3],
        ).permute(0, 1, 3, 4, 2).contiguous()

        predict_x = torch.sigmoid(predict_feature[..., 0])
        predict_y = torch.sigmoid(predict_feature[..., 1])
        predict_w = predict_feature[..., 2]
        predict_h = predict_feature[..., 3]
        predict_obj_conf = torch.sigmoid(predict_feature[..., 4])
        predict_class_conf_list = torch.sigmoid(predict_feature[..., 5:])

        boxes_loss_scale = 2 - boxes_loss_weight_xw * boxes_loss_weight_yh

        loss_x = torch.sum(torch.nn.BCELoss()(predict_x, boxes_x) * boxes_loss_scale * boxes_obj_mask)
        loss_y = torch.sum(torch.nn.BCELoss()(predict_y, boxes_y) * boxes_loss_scale * boxes_obj_mask)

        loss_w = torch.sum(torch.nn.MSELoss()(predict_w, boxes_w) * 0.5 * boxes_loss_scale * boxes_obj_mask)
        loss_h = torch.sum(torch.nn.MSELoss()(predict_h, boxes_h) * 0.5 * boxes_loss_scale * boxes_obj_mask)

        loss_conf = self.lambda_obj * torch.sum(
            torch.nn.BCELoss()(predict_obj_conf, boxes_obj_mask) * boxes_obj_mask) + \
                    self.lambda_noobj * torch.sum(
            torch.nn.BCELoss()(predict_obj_conf, boxes_obj_mask) * boxes_noobj_mask)

        loss_class = torch.sum(torch.nn.BCELoss()(predict_class_conf_list[boxes_obj_mask == 1],
                                                  boxes_class_conf_list[boxes_obj_mask == 1]))

        # print("\n---------------------------------------")
        # print(loss_x, loss_y)
        # print(loss_w, loss_h)
        # print(loss_conf, loss_class)
        # print("---------------------------------------\n")

        loss = loss_x * self.lambda_xy + loss_y * self.lambda_xy + \
               loss_w * self.lambda_wh + loss_h * self.lambda_wh + \
               loss_conf * self.lambda_conf + loss_class * self.lambda_class

        return loss, torch.sum(boxes_obj_mask)

    def forward(self, predict_feature_list,
                tensord_boxes_list: List[torch.Tensor]) -> torch.Tensor:
        boxes_13 = self.decode_pyramid_boxes(tensord_boxes_list, 13, predict_feature_list[0], False)
        boxes_26 = self.decode_pyramid_boxes(tensord_boxes_list, 26, predict_feature_list[1], False)
        boxes_52 = self.decode_pyramid_boxes(tensord_boxes_list, 52, predict_feature_list[2], False)

        loss_13, loss_13_num = self.compute_loss(predict_feature_list[0], boxes_13)
        loss_26, loss_26_num = self.compute_loss(predict_feature_list[1], boxes_26)
        loss_52, loss_52_num = self.compute_loss(predict_feature_list[2], boxes_52)

        loss_list = []

        if not torch.isnan(loss_13):
            loss_list.append(loss_13)
        if not torch.isnan(loss_26):
            loss_list.append(loss_26)
        if not torch.isnan(loss_52):
            loss_list.append(loss_52)

        assert len(loss_list) != 0

        loss = sum(loss_list)

        loss_num = loss_13_num + loss_26_num + loss_52_num

        return loss / loss_num


"""
podpy is an implementatin of the pixel optical depth method as described in
Turner et al. 2014, MNRAS, 445, 794, and Aguirre et al. 2002, ApJ, 576, 1.
Please contact the author (Monica Turner) at turnerm@mit.edu if you have
any questions, comment or issues.
"""

import numpy as np

# Rest wavelength source: Morton 2003, ApJS, 149, 205
# Oscillator strength source: Morton 2003, ApJS, 149, 205
#                                Verner et al 1994, AAPS, 108, 287

# Lyman lines
lambda_h1 = np.array([1215.6701, 1025.7223, 972.5368, 949.7431,
       937.8035, 930.7483, 926.2257, 923.1504,
       920.9631, 919.3514, 918.1294, 917.1806,
       916.429, 915.824, 915.329, 914.919,
       914.576, 914.286, 914.039, 913.826,
       913.641, 913.480, 913.339, 913.215,
       913.104, 913.006, 912.918, 912.839,
       912.768, 912.703, 912.645])
f_h1 = np.array([0.416400, 0.079120, 0.029000, 0.013940,
       0.007799, 0.004814, 0.003183, 0.002216,
       0.001605, 0.00120,  0.000921, 7.226e-4,
       0.000577, 0.000469, 0.000386, 0.000321,
       0.000270, 0.000230, 0.000197, 0.000170,
       0.000148, 0.000129, 0.000114, 0.000101,
       0.000089, 0.000080, 0.000071, 0.000064,
       0.000058, 0.000053, 0.000048])
g_h1 = f_h1 * lambda_h1

lambda_h1_limit = 911.8

# Metals
# OVI
lambda_o6 = np.array([1031.927, 1037.616])
f_o6 = np.array([0.132900, 0.066090])
g_o6 = f_o6 * lambda_o6
# NV
lambda_n5 = np.array([1238.821, 1242.804])
f_n5 = np.array([0.157000, 0.078230])
g_n5 = f_n5 * lambda_n5
# CIV
lambda_c4 = np.array([1548.195, 1550.770])
f_c4 = np.array([0.190800, 0.095220])
g_c4 = f_c4 * lambda_c4
# CIII
lambda_c3 =  np.array([977.020])
f_c3 = np.array([0.7620])
g_c3 = f_c3 * lambda_c3
# CII
lambda_c2 = np.array([1334.5323, 1036.3367])
f_c2 = np.array([ 0.1278, 0.1231])
g_c2 = f_c2 * lambda_c2
# SiIV
lambda_si4 = np.array([1393.755, 1402.770])
f_si4 = np.array([0.5140, 0.2553])
g_si4 = f_si4 * lambda_si4
# SiIII
lambda_si3 = np.array([1206.500])
f_si3 = np.array([1.669000])
g_si3 = f_si3 * lambda_si3
# SiII
lambda_si2 = np.array([1260.4221, 1193.2897, 1190.4158, 989.8731,
	1526.7066, 1304.3702, 1020.6989])
f_si2 = ([1.007000, 0.499100, 0.250200, 0.133000,
	0.11600, 0.09400, 0.028280])
g_si2 = f_si2 * lambda_si2
# MgII
lambda_mg2 = np.array([2796.352, 2803.531])
f_mg2 = np.array([0.6123, 0.3054])
g_mg2 = f_mg2 * lambda_mg2
# OI
lambda_o1 = np.array([1302.1685, 988.7734])
f_o1 = np.array([0.048870, 0.043180])
g_o1 = f_o1 * lambda_o1
# Fe2
lambda_fe2 = np.array([1144.9379, 1608.45085, 1063.1764, 1096.8769,
       1260.533, 1121.9748, 1081.8748, 1143.2260, 1125.4477])
f_fe2 = np.array([0.083, 0.0577, 0.0547, 0.032700,
       0.024000, 0.0290, 0.012600, 0.0192, 0.0156])
g_fe2 = f_fe2 * lambda_fe2

# Constants
c =  299792.458 # speed of light in km/s

# Confidence intervals

from scipy.stats import norm

one_sigma_above = 100. * norm.cdf(1)
two_sigma_above = 100. * norm.cdf(2)
three_sigma_above = 100. * norm.cdf(3)

one_sigma_below = 100. - one_sigma_above
two_sigma_below = 100. - two_sigma_above
three_sigma_below = 100. - three_sigma_above