# -*- coding: utf-8 -*-

#Created on Tue Jun 25 21:59:42 2019
#
#@Author: Zhi-Jiang Yang, Dong-Sheng Cao
#@Institution: CBDD Group, Xiangya School of Pharmaceutical Science, CSU, China，
#@Homepage: http://www.scbdd.com
#@Mail: yzjkid9@gmail.com; oriental-cds@163.com
#@Blog: https://blog.moyule.me

from itertools import combinations
import sys,csv,os
import numpy as np

from rdkit import RDConfig
from rdkit.Chem import AllChem as Chem
from rdkit.Chem import Descriptors, Lipinski, QED
from rdkit.Chem.Scaffolds import MurckoScaffold
sys.path.append(RDConfig.RDContribDir)

from SA_Score import sascorer
from NP_Score import npscorer
from IFG.ifg import identify_functional_groups

try:
    from ..fingerprint import fingerprints
except:
    sys.path.append(__file__+'/..')
    from fingerprint import fingerprints

try:
    from .. import ScoConfig
except:
    sys.path.append('..')
    import ScoConfig

ContriDir = RDConfig.RDContribDir
filename = os.path.join(ContriDir, 'NP_Score/publicnp.model.gz')
fscore = npscorer.pickle.load(npscorer.gzip.open(filename))

def CalculateMolWeight(mol):
    """
    Calculation of molecular weight(contain hydrogen atoms)
    --->MW

    :param mol: molecule
    :type mol: rdkit.Chem.rdchem.Mol
    :return: the weight of molecule(contain hydrogen atoms)
    :rtype: float

    """
    MW = Descriptors.ExactMolWt(mol)
    return round(MW, 2)

def CalculateNumBonds(mol):
    """
    Calculation the number of bonds where between heavy atoms
    --->nBond

    :param mol: molecule
    :type mol: rdkit.Chem.rdchem.Mol
    :return: the number of bonds where between heavy atoms
    :rtype: int

    """
    nBond = mol.GetNumBonds()
    return nBond

def CalculateNumAtoms(mol):
    """
    Calculation of the number of atoms in molecular(contain hydrogen atoms)
    --->nAtom

    :param mol: molecule
    :type mol: rdkit.Chem.rdchem.Mol
    :return: the number of atoms in molecular(contain hydrogen atoms)
    :rtype: int

    """
    mol = Chem.AddHs(mol)
    return mol.GetNumAtoms()

def CalculateNumHetero(mol):
    """
    Calculation of the number of heteroatom in a molecule
    --->nHet

    :param mol: molecule
    :type mol: rdkit.Chem.rdchem.Mol
    :return: the number of heteroatom in a molecule
    :rtype: int

    """
    i = len(
            [atom for atom in mol.GetAtoms()\
             if atom.GetAtomicNum() in [1, 6]]
            )
    nHet = mol.GetNumAtoms()-i
    return nHet

def CalculateNumRotatableBonds(mol):
    """
    Calculation of the number of rotatableBonds
    --->nRot

    Note:
        In some situaion Amide C-N bonds are not considered
        because of their high rotational energy barrier

    :param mol: molecule
    :type mol: rdkit.Chem.rdchem.Mol
    :return: the number of rotatableBond
    :rtype: int

    """
    patt = Chem.MolFromSmarts('[!$(*#*)&!D1]-&!@[!$(*#*)&!D1]')
    nRot = len(mol.GetSubstructMatches(patt))
    return nRot

def CalculateNumRigidBonds(mol):
    """
    Number of non-flexible bonds, in opposite to rotatable bonds
    --->nRig

    :param mol: molecule
    :type mol: rdkit.Chem.rdchem.Mol
    :return: the number of non-flexible bonds
    :rtype: int

    """
    nBOND = CalculateNumBonds(mol)
    flex = 0
    for bond in mol.GetBonds():
        bondtype = bond.GetBondType()
        if bondtype == Chem.rdchem.BondType.SINGLE and not bond.IsInRing():
            flex+=1
    nRig = nBOND-flex
    return nRig

def CalculateFlexibility(mol):
    """
    The flexibility (ration between rotatable and rigid bonds)
    --->Flex

    :param mol: molecule
    :type mol: rdkit.Chem.rdchem.Mol
    :return: the number of ring
    :rtype: float

    """
    nRot = CalculateNumRotatableBonds(mol)
    nRig = CalculateNumRigidBonds(mol)
    return round(nRot/nRig, 2) if nRig else np.float('nan')

def CalculateNumRing(mol):
    """
    Calculation of the number of ring
    --->nRing

    :param mol: molecule
    :type mol: rdkit.Chem.rdchem.Mol
    :return: the number of ring
    :rtype: int

    """
    nRing = Chem.GetSSSR(mol)
    return nRing

def CalculateNumHeavyAtom(mol):
    """
    Calculation of Heavy atom counts in a molecule
    --->nHev

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :return: the number of heavy atom counts in a molecule
    :rtype: int

    """
    nHev = mol.GetNumHeavyAtoms()
    return nHev

def CalculateLogD(mol):
    """
    Calculation of molecular logD under pH=7.4
    --->LogD

    Note:
        We have built a liner model with DNN to predict logD7.4.

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :return: molecular logD under pH=7.4
    :rtype: float

    """
    intercept = 0.5748907159915493

    fps = fingerprints.CalculateGhoseCrippen([mol]).flatten()
    with open(ScoConfig.CrippenDir + '\\Crippen.txt') as f_obj:
        lines = csv.reader(f_obj,delimiter='\t')
        next(lines)
        contri = [x[-1] for x in lines]
        contri = [float(x) for x in contri]
    f_obj.close()
    logD = sum([a*b for a,b in zip(fps,contri)]) + intercept
    return logD

def CalculateLogP(mol):
    """
    Calculation of molecular LogP
    --->logP

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :return: molecular logP
    :rtype: float

    """
    logP = round(Descriptors.MolLogP(mol), 2)
    return logP

def CheckAcid(mol):
    """
    Judge a molecular whether is acid via SMARTS.
    These SMARTS retrived from https://www.daylight.com/dayhtml_tutorials/languages/smarts/smarts_examples.html
    --->ab

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :return: classification to acid or base
    :rtype: str

    """
    acid_fragment = ['[!H0;F,Cl,Br,I,N+,$([OH]-*=[!#6]),+]',
                     '[CX3](=O)[OX2H1]',
                     '[CX3](=O)[OX1H0-,OX2H1]',
                     '[$([OH]-*=[!#6])]',
                     '[$(P(=[OX1])([$([OX2H]),$([OX1-]),$([OX2]P)])([$([OX2H]),$([OX1-]),$([OX2]P)])[$([OX2H]),$([OX1-]),$([OX2]P)]),$([P+]([OX1-])([$([OX2H]),$([OX1-]),$([OX2]P)])([$([OX2H]),$([OX1-]),$([OX2]P)])[$([OX2H]),$([OX1-]),$([OX2]P)])]',
                     '[$([#16X4](=[OX1])(=[OX1])([#6])[OX2H,OX1H0-]),$([#16X4+2]([OX1-])([OX1-])([#6])[OX2H,OX1H0-])]',
                     '[CX3](=[OX1])[F,Cl,Br,I]'
                     ]
    for sma in acid_fragment:
        patt = Chem.MolFromSmarts(sma)
        if mol.HasSubstructMatch(patt):
            return 'acid'
    else:
        return 'base'

def CalculatepKa(mol):
    from math import log10
    """
    Calculating pKa based on the ralation between logD and logP in specific pH.
    --->pKa

    Eq.:
        abs(pH-pKa) = log10(10^(logP-logD)-1)
        pKa = pH - log10(10^(logP-logD)-1) for acid
        pKa = log10(10^(logP-logD)-1) - pH for base

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :return: molecular pKa
    :rtype: float

    """
    logP = CalculateLogP(mol)
    logD = CalculateLogD(mol)
    status = CheckAcid(mol)
    try:
        if status == 'acid':
            pKa = 7.4 - log10(10**(logP-logD)-1)
        else:
            pKa = log10(10**(logP-logD)-1) - 7.4
        return pKa
    except:
        return np.float('nan')

def CalculateMolMR(mol):
    """
    Cacluation of molecular refraction value based on Crippen method
    --->MR

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :return: molecular refraction value based on Crippen method
    :rtype: float

    """
    MR = round(Descriptors.MolMR(mol), 2)
    return MR

def CalculateNumHDonors(mol):
    """
    Caculation of the number of Hydrogen Bond Donors
    --->nHD

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :return: the number of Hydrogen Bond Donors
    :rtype: int

    """
    nHD = Lipinski.NumHDonors(mol)
    return nHD

def CalculateNumHAcceptors(mol):
    """
    Caculation of the number of Hydrogen Bond Acceptors
    --->nHA

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :return: the number of Hydrogen Bond Acceptors
    :rtype: int

    """
    nHA = Lipinski.NumHAcceptors(mol)
    return nHA

def CalculateNumHyBond(mol):
    """
    Sum of Hydrogen Bond Donnors and Acceptors
    --->nHB

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :return: sum of Hydrogen Bond Donnors and Acceptors
    :rtype: int

    """
    nHD = CalculateNumHDonors(mol)
    nHA = CalculateNumHAcceptors(mol)
    nHB = nHD+nHA
    return nHB

def CalculateNumAromaAtom(mol):
    """
    Calculation of aromatic atom counts in a molecule
    --->nAAtom

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :return: the aromatic atom counts in a molecule
    :rtype: int

    """
    aroma = len(mol.GetSubstructMatches(Chem.MolFromSmarts('[a]')))
    return aroma

def CalculateNumAromaRing(mol):
    n = 0
    aatom = mol.GetSubstructMatches(Chem.MolFromSmarts('[a]'))
    aatom = sum(aatom,())
    if aatom:
        ringinfo = mol.GetRingInfo()
        for info in ringinfo.AtomRings():
            n += 1 if set(info) == set(info)&set(aatom) else 0
    else:
        pass
    return n

def CalculateAromaticProportion(mol):
    """
    The proportion of heavy atoms in the molecule that are in an aromatic ring
    --->AP

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :return: the proportion of heavy atoms in the molecule that are in an aromatic ring
    :rtype: float

    """
    aroma = CalculateNumAromaAtom(mol)
    total = CalculateNumHeavyAtom(mol)
    AP = round(aroma/total, 2) if total else np.float('nan')
    return AP

def CalculateLogSw(mol):
    """
    The logSw represents the logarithm of compounds water solubility computed by the ESOL method
    --->logSw

    Equation:
        Log(Sw) = 0.16-0.638*clogP-0.0062*MWT+0.066*RB-0.74*AP
        where, MWT: Molecular Weight; RB: Rotatable bonds; AP: Aromatic proportion

    Reference:
        (1) `Delaney, John S (2004)`_.

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :return: the molecular logSw
    :rtype: float

    .. _Delaney, John S (2004):
        https://pubs.acs.org/doi/abs/10.1021/ci034243x

    """
    #Calculate each property
    MWT = CalculateMolWeight(mol)
    RB = CalculateNumRotatableBonds(mol)
    AP = CalculateAromaticProportion(mol)
    logP = CalculateLogP(mol)
    logSw = 0.16-0.638*logP-0.0062*MWT+0.066*RB-0.74*AP
    return round(logSw,2)

def CalculateFsp3(mol):
    """
    Fsp3 (carbon bond saturation) is defined as the number of sp3 hybridized carbons / total carbon count.
    --->FSP3

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :return: the carbon bond saturation
    :rtype: float

    """
    return round(Lipinski.FractionCSP3(mol), 2)

def CalculateTPSA(mol):
    """
    Calculation of TPSA
    --->TPSA

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :return: TPSA
    :rtype: float

    """
    TPSA = round(Descriptors.TPSA(mol), 2)
    return TPSA

def CalculateQEDmean(mol):
    """
    Calculation QED descriptor under different weights
    A descriptor a measure of drug-likeness based on the concept of desirability
    Here, calculating the QED descriptor using average descriptor weights.
    --->QEDmean

    Reference:
        (1) `Bickerton, G. Richard (2012)`_.

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :return: QED descriptor using average descriptor weights
    :rtype: float

    .. _Bickerton, G. Richard (2012):
        https://www.nature.com/nchem/journal/v4/n2/abs/nchem.1243.html

    """
    QEDmean = QED.weights_mean(mol)
    return round(QEDmean, 2)

def CalculateQEDmax(mol):
    """
    Calculation QED descriptor under different weights
    A descriptor a measure of drug-likeness based on the concept of desirability
    Here, calculating the QED descriptor using maximal descriptor weights.
    --->QEDmax

    Reference:
        (1) `Bickerton, G. Richard (2012)`_.

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :return: QED descriptor using maximal descriptor weights
    :rtype: float

    .. _Bickerton, G. Richard (2012):
        https://www.nature.com/nchem/journal/v4/n2/abs/nchem.1243.html

    """
    QEDmax = QED.weights_max(mol)
    return round(QEDmax,2)

def CalculateQEDnone(mol):
    """
    Calculation QED descriptor under different weights
    A descriptor a measure of drug-likeness based on the concept of desirability
    Here, calculating the QED descriptor using unit weights.
    --->QEDnone

    Reference:
        (1) `Bickerton, G. Richard (2012)`_.

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :return: QED descriptor using unit weights
    :rtype: float

    .. _Bickerton, G. Richard (2012):
        https://www.nature.com/nchem/journal/v4/n2/abs/nchem.1243.html

    """
    QEDnone = QED.weights_none(mol)
    return round(QEDnone,2)

def CalculateMaxSizeSystemRing(mol):
    """
    Number of atoms involved in the biggest system ring
    ---> maxring

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :return: number of atoms involved in the biggest system ring
    :rtype: int

    """
    #0.Get the scaffold
    core = MurckoScaffold.GetScaffoldForMol(mol)
    fw = MurckoScaffold.MakeScaffoldGeneric(core)
    #1.Obtaining which atoms consist of rings
    MaxRing = 0
    ri = fw.GetRingInfo()
    atoms = list(ri.AtomRings())
    length = len(atoms)
    if length == 0:
        pass
    else:
        rw = Chem.RWMol(fw)
        #2.Judge which atoms are replacement
        atoms = [set(x) for x in atoms]
        for pair in combinations(range(length),2):
            replace = list(atoms[pair[0]]&atoms[pair[1]])
            if len(replace) >= 2:
                for repl in list(combinations(replace,2)):
                    rw.RemoveBond(*repl)
            else:
                pass
        m = Chem.MolFromSmiles(Chem.MolToSmiles(rw))
        ri = m.GetRingInfo()
        bonds = ri.BondRings()
        for item in bonds:
            if len(item) > MaxRing:
                MaxRing = len(item)
    return MaxRing

def CalculateNumStereocenters(mol):
    """
    the number of stereo centers
    --->nStereo

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :return: the number of stereo centers
    :rtype: int

    """
    return Chem.CalcNumAtomStereoCenters(mol)

def _CalculateNumElement(mol, AtomicNumber=6):
    """
    **Internal used only**
    Calculation of specific type of atom number in a molecule

    Calculation of element counts with atomic number equal to n in a molecule

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :param AtomicNumber: the AtomicNumber of atom to be counted, defaults to 6
    :type AtomicNumber: int, optional
    :return: the number of stereo centers
    :rtype: int

    """
    return len(
            [atom for atom in mol.GetAtoms()\
                if atom.GetAtomicNum() == AtomicNumber]
            )

def CalculateNumCarbon(mol):
    """
    Calculation of Carbon number in a molecule
    --->nC

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :return: the number of carbon atoms
    :rtype: int

    """
    return _CalculateNumElement(mol,AtomicNumber=6)

def CalculateNumBoron(mol):
    """
    Calculation of Boron counts in a molecule
    --->nB

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :return: the number of boron atoms
    :rtype: int

    """
    return _CalculateNumElement(mol,AtomicNumber=5)

def CalculateNumFluorin(mol):
    """
    Calculation of Fluorin counts in a molecule
    --->nF

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :return: the number of fluori atoms
    :rtype: int

    """
    return _CalculateNumElement(mol, AtomicNumber=10)

def CalculateNumChlorin(mol):
    """
    Calculation of Chlorin counts in a molecule
    --->nCl

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :return: the number of chlorin atoms
    :rtype: int

    """
    return _CalculateNumElement(mol, AtomicNumber=17)

def CalculateNumBromine(mol):
    """
    Calculation of Bromine counts in a molecule
    --->nBr

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :return: the number of bromine atoms
    :rtype: int

    """
    return _CalculateNumElement(mol, AtomicNumber=35)

def CalculateNumIodine(mol):
    """
    Calculation of Iodine counts in a molecule
    --->nI

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :return: the number of bromine atoms
    :rtype: int

    """
    return _CalculateNumElement(mol, AtomicNumber=53)

def CalculateNumPhosphor(mol):
    """
    Calcualtion of Phosphor number in a molecule
    --->nP

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :return: the number of phosphor atoms
    :rtype: int

    """
    return _CalculateNumElement(mol,AtomicNumber=15)

def CalculateNumSulfur(mol):
    """
    Calculation of Sulfur counts in a molecule
    --->nS

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :return: the number of sulfur atoms
    :rtype: int

    """
    return _CalculateNumElement(mol,AtomicNumber=16)

def CalculateNumOxygen(mol):
    """
    Calculation of Oxygen counts in a molecule
    --->nO

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :return: the number of oxygen atoms
    :rtype: int

    """
    return _CalculateNumElement(mol,AtomicNumber=8)

def CalculateNumNitrogen(mol):
    """
    Calculation of Nitrogen counts in a molecule
    --->nN

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :return: the number of nitrogen atoms
    :rtype: int

    """
    return _CalculateNumElement(mol,AtomicNumber=7)

# def CalculateNumberChargedGroups(mol):
#     """
#     Number of Charged Groups
#     --->nChar

#     :param mol: molecular
#     :type mol: rdkit.Chem.rdchem.Mol
#     :return: the number of charged group
#     :rtype: int

#     """
#     pass

def CalculateHetCarbonRatio(mol):
    """
    The ratio between the number of non carbon atoms and the number of carbon atoms.
    --->HetRatio

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :return: the ratio between the number of non carbon atoms and the number of carbon atoms
    :rtype: float

    """
    nHet = CalculateNumHetero(mol)
    nCarb = CalculateNumCarbon(mol)
    return round(nHet/nCarb,2) if nCarb else np.float('nan')

def CalculateSAscore(mol):
    """
    A function to estimate ease of synthesis (synthetic accessibility) of drug-like molecules
    --->SAscore

    Reference:
        (1) `Ertl Peter (2009)`_.

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :return: ease of synthesis
    :rtype: float

    .. _Ertl Peter (2009):
        https://jcheminf.biomedcentral.com/articles/10.1186/1758-2946-1-8

    """
    return round(sascorer.calculateScore(mol), 2)

def CalculateNPscore(mol):
    """
    A function to calculate the natural product-likeness score
    --->NPscore

    Reference:
        (1) `Ertl Peter (2008)`_.

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :return: product-likeness score
    :rtype: float

    .. _Ertl Peter (2008):
        https://pubs.acs.org/doi/abs/10.1021/ci700286x

    """
    return round(npscorer.scoreMol(mol,fscore=fscore), 2)

def GetIFG(mol):
    """
    A function to compute functional groups in organic molecules
    --->IFG

    Reference:
        (1) `Ertl Peter (2017)`_.

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :return: list of namedtuple, namedtuple('IFG', ['atomIds', 'atoms', 'type'])
    :rtype: list

    .. _Ertl Peter (2017):
        https://jcheminf.biomedcentral.com/articles/10.1186/s13321-017-0225-z

    """
    return [fg._asdict() for fg in identify_functional_groups(mol)]

def CalculateMolVolume(mol):
    """
    Calculation of Van der Waals Volume of molecule
    --->MV

    Equation:
        for single atom: Vw = 4/3*pi*rw^3, the rw is the Van der Waals radius of atom
        VvdW = ∑(atom contributions)-5.92NB(Unit in Å^3), NB is the total number of bonds
        the Van der Waals radius of atom is derived from wikipedia.

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :return: Van der Waals Volume of molecule
    :rtype: float

    """
    from math import pi
    Radii = {'H':1.20,'C':1.70,'N':1.55,
             'O':1.52,'S':1.80,'P':1.80,
             'F':1.47,'Cl':1.75,'Br':1.85,
             'I':1.98,'Na':2.27,'Mg':1.73,
             'K':2.75,'Ca':2.31,'Ba':2.68,
             'He':140,'Li':182,'Be':153,
             'B':192,'Ne':154,'Al':184,
             'Si':210,'Ar':188,'Ni':163,
             'Cu':140,'Zn':139,'Ga':187,
             'Ge':211,'As':185,'Se':190,
             'Kr':202,'Rb':303,'Sr':249,
             'Pd':163,'Ag':172,'Cd':158,
             'In':193,'Sn':217,'Sb':206,
             'Te':206,'Xe':216,'Cs':343,
             'Pt':175,'Au':166,'U':186,
             'Hg':155,'Tl':196,'Pb':202,
             'Bi':207,'Po':197,'At':202,
             'Rn':220,'Fr':348,'Ra':283}
    mol = Chem.AddHs(mol)
    contrib = []
    for atom in mol.GetAtoms():
        try:
            contrib.append(Radii[atom.GetSymbol()])
        except:
            pass
    # contrib = [Radii[atom.GetSymbol()] for atom in mol.GetAtoms()]
    contrib = [pi*(r**3)*4/3 for r in contrib]
    vol = sum(contrib) - 5.92*len(mol.GetBonds())
    return round(vol, 2)

def CalculateMolDensity(mol):
    """
    Calculation of density of molecule
    --->Dense

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :return: density of molecule
    :rtype: float

    """
    MW = CalculateMolWeight(mol)
    Vol = CalculateMolVolume(mol)
    return round(MW/Vol, 2) if Vol else np.float('nan')

def CalculateMolFCharge(mol):
    """
    Calculation of formal charge of molecule
    --->fChar

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :return: formal charge of molecule
    :rtype: float

    """
    mol = Chem.AddHs(mol)
    FChar = [atom.GetFormalCharge() for atom in mol.GetAtoms()]
    return sum(FChar)

def _CalculateNumBond(mol,btype):
    """
    **Internal used only**
    Calculation of specific type of bond number in a molecule

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :param btype: the type of bond to be counted
    :type btype: str
    :return: the number of specific bond
    :rtype: int

    """
    if btype == 'SINGLE':
        return len([bond for bond in mol.GetBonds()
                    if bond.GetBondType() == Chem.rdchem.BondType.SINGLE])
    elif btype == 'DOUBLE':
        return len([bond for bond in mol.GetBonds()
                    if bond.GetBondType() == Chem.rdchem.BondType.DOUBLE])
    elif btype == 'TRIPLE':
        return len([bond for bond in mol.GetBonds()
                    if bond.GetBondType() == Chem.rdchem.BondType.TRIPLE])

def CalculateNumSinBond(mol):
    """
    Calculation of single bond number of molecule
    ---> nSingle

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :return: the number of single bond
    :rtype: int

    """
    return _CalculateNumBond(mol, btype='SINGLE')

def CalculateNumDouBond(mol):
    """
    Calculation of double bond number of molecule
    --->nDouble

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :return: the number of double bond
    :rtype: int

    """
    return _CalculateNumBond(mol, btype='DOUBLE')

def CalculateNumTriBond(mol, btype='TRIPLE'):
    """
    Calculation of triple bond number of molecule
    ---> nTriple

    :param mol: molecular
    :type mol: rdkit.Chem.rdchem.Mol
    :return: the number of triple bond
    :rtype: int

    """
    return _CalculateNumBond(mol, btype='TRIPLE')

def GetProperties(mol,
                  items=['MW','Vol','Dense','fChar','nBond','nAtom','nHD','nHA','nHB',
                         'nHet','nStero','nHev','nRot','nRig','Flex','nRing','logP',
                         'logD','pKa','logSw','ab','MR','TPSA','AP','HetRatio','Fsp3',
                         'MaxRing','QEDmean','QEDmax','QEDnone','SAscore','NPscore',
                         'nSingle','nDouble','nTriple','nC','nB','nF','nCl','nBr','nI',
                         'nP','nS','nO','nN']
        ):
    """
    Get all properties in scopy
    """
    funcl = {'MW': 'CalculateMolWeight(mol)',
    'Vol': 'CalculateMolVolume(mol)',
    'Dense': 'CalculateMolDensity(mol)',
    'fChar': 'CalculateMolFCharge(mol)',
    'nBond': 'CalculateNumBonds(mol)',
    'nAtom': 'CalculateNumAtoms(mol)',
    'nHet': 'CalculateNumHetero(mol)',
    'nRot': 'CalculateNumRotatableBonds(mol)',
    'nRig': 'CalculateNumRigidBonds(mol)',
    'Flex': 'CalculateFlexibility(mol)',
    'nRing': 'CalculateNumRing(mol)',
    'nHev': 'CalculateNumHeavyAtom(mol)',
    'logP': 'CalculateLogP(mol)',
    'logD': 'CalculateLogD(mol)',
    'pKa': 'CalculatepKa(mol)',
    'ab': 'CheckAcid(mol)',
    'MR': 'CalculateMolMR(mol)',
    'nHD': 'CalculateNumHDonors(mol)',
    'nHA': 'CalculateNumHAcceptors(mol)',
    'nHB': 'CalculateNumHyBond(mol)',
    'AP': 'CalculateAromaticProportion(mol)',
    'logSw': 'CalculateLogSw(mol)',
    'Fsp3': 'CalculateFsp3(mol)',
    'TPSA': 'CalculateTPSA(mol)',
    'MaxRing': 'CalculateMaxSizeSystemRing(mol)',
    'nStero': 'CalculateNumStereocenters(mol)',
    'HetRatio': 'CalculateHetCarbonRatio(mol)',
    'QEDmean': 'CalculateQEDmean(mol)',
    'QEDmax': 'CalculateQEDmax(mol)',
    'QEDnone': 'CalculateQEDnone(mol)',
    'SAscore': 'CalculateSAscore(mol)',
    'NPscore': 'CalculateNPscore(mol)',
    'nSingle': 'CalculateNumSinBond(mol)',
    'nDouble': 'CalculateNumDouBond(mol)',
    'nTriple': 'CalculateNumTriBond(mol)',
    'nC': 'CalculateNumCarbon(mol)',
    'nB': 'CalculateNumBoron(mol)',
    'nF': 'CalculateNumFluorin(mol)',
    'nCl': 'CalculateNumChlorin(mol)',
    'nBr': 'CalculateNumBromine(mol)',
    'nI': 'CalculateNumIodine(mol)',
    'nP': 'CalculateNumPhosphor(mol)',
    'nS': 'CalculateNumSulfur(mol)',
    'nO': 'CalculateNumOxygen(mol)',
    'nN': 'CalculateNumNitrogen(mol)'}

    vals = []
    for item in items:
        val = eval(funcl[item])
        vals.append(val)

    return dict(zip(items, vals))

if __name__ =='__main__':

    smis = [
            'C1=CC=CC(C(Br)C)=C1',
            'C1=CC2NC(=O)CC3C=2C(C(=O)C2C=CC=CC=23)=C1',
            'C1=CC=C2C(=O)C3C=CNC=3C(=O)C2=C1',
            'C1=NC(CCN)=CN1',
            'C1CCCC(CCO)C1',
            'C1=CC=C2N=C(O)C=CC2=C1',
            'C(OC)1=C(C)C=C2OC[C@]([H])3OC4C(C)=C(OC)C=CC=4C(=O)[C@@]3([H])C2=C1C',
            'C1=C2N=CC=NC2=C2N=CNC2=C1',
            'C1=C(O)C=CC(O)=C1',
            'CCC1(c2ccccc2)C(=O)NC(=O)NC1=O',
            'N1=CN=CN=C1',
            'C1=C2C=CC=CC2=CC2C=CC=CC1=2', #NonGenotoxic_Carcinogenicity
            'C1=CC=C2C(=O)CC(=O)C2=C1', #Pains
            'C1=CC=CC(COCO)=C1', #Potential_Electrophilic
            'N1=NC=CN1C=O', #Promiscuity
            'CC(=O)OC(=O)C1C=COC1', #Skin_Sensitization
            'S',
            'CCCCC(=O)[H]', #Biodegradable
            'C1=CN=C(C(=O)O)C=C1', #Chelating
            'C(OC)1=CC=C2OCC3OC4C=C(OC)C=CC=4C(=O)C3C2=C1',
            'C1=C2N=CC=NC2=C2N=CNC2=C1', #Genotoxic_Carcinogenicity_Mutagenicity
            'N(CC)(CCCCC)C(=S)N', #Idiosyncratic
            ]

    for index, smi in enumerate(smis):
        mol = Chem.MolFromSmiles(smi)
        print('Index:{}'.format(index))
        res = GetProperties(mol=mol)
        print(res)


""" This code implements Q-learning from a static batch of data with KL-control,
monte-carlo target Q-value estimation, Psi-learning, and model averaging."""
import sys
import os
import copy
import pickle
import numpy as np
import pandas as pd
from collections import namedtuple
from itertools import count
import random
from math import isnan
from pathlib import Path
from datetime import datetime

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.autograd as autograd
import torch.optim as optim

from model.solver import Solver, VariationalSolver
from model.models import VariationalModels
from model.utils import to_var, pad, EOS_ID, TensorboardWriter, convert_old_checkpoint_format
from model.data_loader import get_loader
import replay_buffer
import rewards

class BatchQ:
    def __init__(self, config, val_config):
        self.config = config
        self.val_config = val_config

        # Load experience replay buffer from file
        self.experience = replay_buffer.CsvReplayBuffer(
            config.experience_path, raw=config.raw_buffer,
            history_len=config.max_conversation_length, config=config,
            max_sentence_length=config.max_sentence_length,
            rewards=config.rewards, reward_weights=config.reward_weights,
            model_averaging=config.model_averaging)
        self.vocab = self.experience.vocab
        self.config.vocab_size = self.experience.vocab.vocab_size
        self.action_dim = self.experience.vocab.vocab_size

        # Check that all required rewards are in the buffer; if not, compute
        for r in config.rewards:
            if r not in self.experience.buffer.columns.values:
                reward_func = getattr(rewards, r)
                self.experience = reward_func(self.experience)

        # Build internal hierarchical models
        self.eval_data = self.get_data_loader()
        self.build_models()

        if self.config.load_rl_ckpt:
            self.load_models()

        self.q_optimizer = torch.optim.Adam(
            filter(lambda p: p.requires_grad, self.q_net.model.parameters()),
            lr=self.config.learning_rate)

        self.set_up_logging()

    def q_update(self):
        """General Q learning update."""
        # Sample a batch
        batch = self.experience.sample(self.config.rl_batch_size)

        # Run underlying q network to get q value of each word in each
        # conversation in the batch. Use the same data to run the prior network
        # and get the rewards based on KL divergence from the prior.
        q_values, prior_rewards = self.get_q_values(batch)

        # Compute target Q values. These will include the rewards observed in
        # the batch (i.e. r + done * gamma * max_a' Q_T(a,s'))
        with torch.no_grad():
            target_q_values = self.get_target_q_values(batch, prior_rewards)

        loss_func = getattr(F, self.config.q_loss_func)
        loss = loss_func(q_values, target_q_values)

        assert not isnan(loss.item())
        self.q_loss_batch_history.append(loss.item())

        # Optimize the model
        self.q_optimizer.zero_grad()
        loss.backward()

        # Clip gradients - absolutely crucial
        torch.nn.utils.clip_grad_value_(self.q_net.model.parameters(),
                                        self.config.gradient_clip)

        self.q_optimizer.step()

        # Update Target Networks
        tau = self.config.target_update_rate
        for param, target_param in zip(self.q_net.model.parameters(),
                                       self.target_q_net.model.parameters()):
            target_param.data.copy_(
                tau * param.data + (1 - tau) * target_param.data)

    def get_q_values(self, batch):
        """update where states are whole conversations which
        each have several sentences, and actions are a sentence (series of
        words). Q values are per word. Target Q values are over the next word
        in the sentence, or, if at the end of the sentence, the first word in a
        new sentence after the user response.
        """
        actions = to_var(torch.LongTensor(batch['action']))  # [batch_size]

        # Prepare inputs to Q network
        conversations = [np.concatenate(
            (conv, np.atleast_2d(batch['action'][i])))
            for i, conv in enumerate(batch['state'])]
        sent_lens = [np.concatenate(
            (lens, np.atleast_1d(batch['action_lens'][i])))
            for i, lens in enumerate(batch['state_lens'])]
        target_conversations = [conv[1:] for conv in conversations]
        conv_lens = [len(c) - 1 for c in conversations]
        if self.config.model not in VariationalModels:
            conversations = [conv[:-1] for conv in conversations]
            sent_lens = np.concatenate([l[:-1] for l in sent_lens])
        else:
            sent_lens = np.concatenate([l for l in sent_lens])
        conv_lens = to_var(torch.LongTensor(conv_lens))

        # Run Q network. Will produce [num_sentences, max sent len, vocab size]
        all_q_values = self.run_seq2seq_model(
            self.q_net, conversations, sent_lens, target_conversations,
            conv_lens)

        # Index to get only q values for actions taken (last sentence in each
        # conversation)
        start_q = torch.cumsum(torch.cat(
            (to_var(conv_lens.data.new(1).zero_()), conv_lens[:-1])), 0)
        conv_q_values = torch.stack(
            [all_q_values[s+l-1, :, :]
             for s, l in zip(start_q.data.tolist(), conv_lens.data.tolist())],
             0)  # [num_sentences, max_sent_len, vocab_size]

        # Limit by actual sentence length (remove padding) and flatten into
        # long list of words
        word_q_values = torch.cat(
            [conv_q_values[i, :l, :]
             for i, l in enumerate(batch['action_lens'])],
             0)  # [total words, vocab_size]
        word_actions = torch.cat(
            [actions[i, :l] for i, l in enumerate(batch['action_lens'])],
            0)  # [total words]

        # Extract q values corresponding to actions taken
        q_values = word_q_values.gather(
            1, word_actions.unsqueeze(1)).squeeze()  # [total words]

        """ Compute KL metrics """
        prior_rewards = None

        # Get probabilities from policy network
        q_dists = torch.nn.functional.softmax(word_q_values, 1)
        q_probs = q_dists.gather(
                1, word_actions.unsqueeze(1)).squeeze()

        with torch.no_grad():
            # Run pretrained prior network.
            # [num_sentences, max sent len, vocab size]
            all_prior_logits = self.run_seq2seq_model(
                self.pretrained_prior, conversations, sent_lens,
                target_conversations, conv_lens)

            # Get relevant actions. [num_sentences, max_sent_len, vocab_size]
            conv_prior = torch.stack(
                [all_prior_logits[s+l-1, :, :]
                for s, l in zip(
                    start_q.data.tolist(), conv_lens.data.tolist())], 0)

            # Limit by actual sentence length (remove padding) and flatten.
            # [total words, vocab_size]
            word_prior_logits = torch.cat(
                [conv_prior[i, :l, :]
                for i, l in enumerate(batch['action_lens'])], 0)

            # Take the softmax
            prior_dists = torch.nn.functional.softmax(
                word_prior_logits, 1)

            kl_div = F.kl_div(q_dists.log(), prior_dists, reduce=False)

            # [total words]
            prior_probs = prior_dists.gather(
                1, word_actions.unsqueeze(1)).squeeze()
            logp_logq = prior_probs.log() - q_probs.log()

            if self.config.model_averaging:
                model_avg_sentences = batch['model_averaged_probs']

                # Convert to tensors and flatten into [num_words]
                word_model_avg = torch.cat([to_var(
                    torch.FloatTensor(m)) for m in model_avg_sentences], 0)

                # Compute KL from model-averaged prior
                prior_rewards = word_model_avg.log() - q_probs.log()

                # Clip because KL should never be negative, so because we
                # are subtracting KL, rewards should never be positive
                prior_rewards = torch.clamp(prior_rewards, max=0.0)

            elif self.config.kl_control and self.config.kl_calc == 'integral':
                # Note: we reward the negative KL divergence to ensure the
                # RL model stays close to the prior
                prior_rewards = -1.0 * torch.sum(kl_div, dim=1)
            elif self.config.kl_control:
                prior_rewards = logp_logq

            if self.config.kl_control:
                prior_rewards = prior_rewards * self.config.kl_weight_c
                self.kl_reward_batch_history.append(
                    torch.sum(prior_rewards).item())

            # Track all metrics
            self.kl_div_batch_history.append(torch.mean(kl_div).item())
            self.logp_batch_history.append(torch.mean(prior_probs.log()).item())
            self.logp_logq_batch_history.append(torch.mean(logp_logq).item())

        return q_values, prior_rewards

    def get_target_q_values(self, batch, prior_rewards=None):
        rewards = to_var(torch.FloatTensor(batch['rewards']))  # [batch_size]
        not_done = to_var(torch.FloatTensor(1 - batch['done']))  # [batch_size]
        self.sampled_reward_batch_history.append(torch.sum(rewards).item())

        # Prepare inputs to target Q network. Append a blank sentence to get
        # best response at next utterance to user input. (Next state
        # includes user input).
        blank_sentence = np.zeros((1, self.config.max_sentence_length))
        next_state_convs = [np.concatenate(
            (conv, blank_sentence)) for conv in batch['next_state']]
        next_state_lens = [np.concatenate(
            (lens, [1])) for lens in batch['next_state_lens']]
        next_targets = [conv[1:] for conv in next_state_convs]
        next_conv_lens = [len(c) - 1 for c in next_state_convs]
        if self.config.model not in VariationalModels:
            next_state_convs = [conv[:-1] for conv in next_state_convs]
            next_state_lens = np.concatenate(
                [l[:-1] for l in next_state_lens])
        else:
            next_state_lens = np.concatenate([l for l in next_state_lens])
        next_conv_lens = to_var(torch.LongTensor(next_conv_lens))

        # [monte_carlo_count, num_sentences, max sent len, vocab size]
        _mc_target_q_values = [[]] * self.config.monte_carlo_count
        for t in range(self.config.monte_carlo_count):
            # Run target Q network. Output is size:
            # [num_sentences, max sent len, vocab size]
            if self.config.monte_carlo_count == 1:
                # In this setting, we don't use dropout out at inference time at all
                all_target_q_values = self.run_seq2seq_model(
                    self.target_q_net, next_state_convs, next_state_lens,
                    next_targets, next_conv_lens)
            else:
                # In this setting, each time we draw a new dropout mask (at inference time)
                all_target_q_values = self.run_seq2seq_model(
                    self.target_q_net, next_state_convs, next_state_lens,
                    next_targets, next_conv_lens)

            # Target indexing: last sentence is a blank to get value of next
            # response. Second last is the user response. 3rd last is models own
            # actions. Note that targets begin at the 2nd word of each sentence.
            start_t = torch.cumsum(torch.cat(
                (to_var(next_conv_lens.data.new(1).zero_()),
                    next_conv_lens[:-1])), 0)
            conv_target_q_values = torch.stack(
                [all_target_q_values[s+l-3, 1:, :]
                for s, l in zip(start_t.data.tolist(),
                                next_conv_lens.data.tolist())],
                0)  # Dimension [num_sentences, max_sent_len - 1, vocab_size]

            # At the end of a sentence, want value of starting a new response
            # after user's response. So index into first word of last blank
            # sentence that was appended to the end of the conversation.
            next_response_targets = torch.stack(
                [all_target_q_values[s+l-1, 0, :]
                for s, l in zip(start_t.data.tolist(),
                                next_conv_lens.data.tolist())], 0)
            next_response_targets = torch.reshape(
                next_response_targets,
                [self.config.rl_batch_size, 1, -1]
                ) # [num_sentences, 1, vocab_size]
            conv_target_q_values = torch.cat(
                [conv_target_q_values, next_response_targets],
                1)  # [num_sentences, max_sent_len, vocab_size]

            # Limit target Q values by conversation length
            limit_conv_targets = [conv_target_q_values[i, :l, :]
                for i, l in enumerate(batch['action_lens'])]

            if self.config.psi_learning:
                # Target is r + gamma * log sum_a' exp(Q_target(s', a'))
                conv_max_targets = [torch.distributions.utils.log_sum_exp(c)
                    for c in limit_conv_targets]
                target_q_values = torch.cat(
                    [rewards[i] + not_done[i] * self.config.gamma * c.squeeze()
                     for i, c in enumerate(conv_max_targets)], 0)  # [total words]
            else:
                # Target is r + gamma * max_a' Q_target(s',a'). Reward and done are
                # at the level of conversation, so add and multiply in before
                # flattening and taking max.
                word_target_q_values = torch.cat(
                    [rewards[i] + not_done[i] * self.config.gamma * c
                     for i, c in enumerate(limit_conv_targets)],
                    0)  # [total words, vocab_size]
                target_q_values, _ = word_target_q_values.max(1)

            _mc_target_q_values[t] = target_q_values
        mc_target_q_values = torch.stack(_mc_target_q_values, 0)

        min_target_q_values, _ = mc_target_q_values.min(0)

        if self.config.kl_control:
            min_target_q_values += prior_rewards

        return min_target_q_values

    def q_learn(self):
        self.q_loss_history = []
        self.q_loss_batch_history = []
        self.sampled_reward_history = []
        self.sampled_reward_batch_history = []
        if self.config.kl_control:
            self.kl_reward_history = []
            self.kl_reward_batch_history = []

        # Need to track KL metrics even for baselines for plots
        self.kl_div_history = []
        self.kl_div_batch_history = []
        self.logp_history = []
        self.logp_batch_history = []
        self.logp_logq_history = []
        self.logp_logq_batch_history = []

        print('Commencing training at step', self.t)
        while self.t <= self.config.num_steps:
            self.q_update()

            # Log metrics
            if self.t % self.config.log_every_n == 0:
                self.epoch_q_loss = np.sum(self.q_loss_batch_history) \
                    / self.config.log_every_n
                self.q_loss_history.append(self.epoch_q_loss)
                self.q_loss_batch_history = []
                print('Average Q loss at step', self.t, '=', self.epoch_q_loss)

                self.epoch_sampled_reward = np.sum(
                    self.sampled_reward_batch_history) / self.config.log_every_n
                self.sampled_reward_history.append(self.epoch_sampled_reward)
                self.sampled_reward_batch_history = []
                print('\tAverage sampled batch reward =',
                    self.epoch_sampled_reward)

                if self.config.kl_control:
                    self.epoch_kl_reward = np.sum(
                        self.kl_reward_batch_history) \
                        / self.config.log_every_n
                    self.kl_reward_history.append(
                        self.epoch_kl_reward)
                    self.kl_reward_batch_history = []
                    print('\tAverage data prior reward =',
                        self.epoch_kl_reward)

                # Logging KL for plots
                self.epoch_kl_div = np.sum(
                    self.kl_div_batch_history) / self.config.log_every_n
                self.kl_div_history.append(self.epoch_kl_div)
                self.kl_div_batch_history = []
                self.epoch_logp = np.sum(
                    self.logp_batch_history) / self.config.log_every_n
                self.logp_history.append(self.epoch_logp)
                self.logp_batch_history = []
                self.epoch_logp_logq = np.sum(
                    self.logp_logq_batch_history) / self.config.log_every_n
                self.logp_logq_history.append(self.epoch_logp_logq)
                self.logp_logq_batch_history = []

                sys.stdout.flush()
                self.write_summary(self.t)

            if self.t > 0 and self.t % self.config.save_every_n == 0:
                self.save_model(self.t)

            self.t += 1

    def build_models(self):
        config = copy.deepcopy(self.config)

        # If loading RL checkpoint, ensure it doesn't try to load the ckpt
        # through Solver
        if self.config.load_rl_ckpt:
            config.checkpoint = None

        if self.config.model in VariationalModels:
            self.q_net = VariationalSolver(
                config, None, self.eval_data, vocab=self.vocab, is_train=True)
            self.target_q_net = VariationalSolver(
                config, None, self.eval_data, vocab=self.vocab, is_train=True)
        else:
            self.q_net = Solver(
                config, None, self.eval_data, vocab=self.vocab, is_train=True)
            self.target_q_net = Solver(
                config, None, self.eval_data, vocab=self.vocab, is_train=True)
        print('Building Q network')
        self.q_net.build()

        print('\nBuilding Target Q network')
        self.target_q_net.build()

        if self.config.model in VariationalModels:
            self.pretrained_prior = VariationalSolver(
                self.config, None, self.eval_data, vocab=self.vocab,
                is_train=True)
        else:
            self.pretrained_prior = Solver(
                self.config, None, self.eval_data, vocab=self.vocab,
                is_train=True)
        print('Building prior network')
        self.pretrained_prior.build()

        # Freeze the weights of the prior so it stays constant
        self.pretrained_prior.model.eval()
        for params in self.pretrained_prior.model.parameters():
            params.requires_grad = False

        print('Successfully initialized Q networks')
        self.t = 0

    def run_seq2seq_model(self, q_net, input_conversations, sent_lens,
                     target_conversations, conv_lens):
        # Prepare the batch
        sentences = [sent for conv in input_conversations for sent in conv]
        targets = [sent for conv in target_conversations for sent in conv]

        if not (np.all(np.isfinite(sentences))
                and np.all(np.isfinite(targets))
                and np.all(np.isfinite(sent_lens))):
            print("Input isn't finite")

        sentences = to_var(torch.LongTensor(sentences))
        targets = to_var(torch.LongTensor(targets))
        sent_lens = to_var(torch.LongTensor(sent_lens))

        # Run Q network
        q_outputs = q_net.model(sentences, sent_lens, conv_lens, targets,
                                rl_mode=True)
        return q_outputs[0]  # [num_sentences, max_sentence_len, vocab_size]

    def write_summary(self, t):
        metrics_to_log = ['epoch_q_loss', 'epoch_sampled_reward',
                          'epoch_kl_div', 'epoch_logp', 'epoch_logp_logq']

        if self.config.kl_control:
            metrics_to_log.append('epoch_kl_reward')

        metrics_dict = {}
        for metric in metrics_to_log:
            met_val = getattr(self, metric, None)
            metrics_dict[metric] = met_val
            if met_val is not None:
                self.writer.update_loss(
                    loss=met_val,
                    step_i=t,
                    name=metric)

        # Write pandas csv with metrics to save dir
        self.df = self.df.append(metrics_dict, ignore_index=True)
        self.df.to_csv(self.pandas_path)

    def set_up_logging(self):
        # Get save path
        time_now = datetime.now().strftime('%Y-%m-%d_%H:%M:%S')
        default_save_path = Path('model_checkpoints/rl/')

        # Folder for type of RL used
        experiment_name = self.config.experiment_name
        if experiment_name is None:
            # if self.config.double_q: experiment_name = 'double_q'
            if self.config.model_averaging:
                experiment_name = 'model_averaging'
            elif self.config.kl_control:
                experiment_name = 'kl_control'
                if self.config.kl_calc == 'sample': experiment_name += '_sample'
            else: experiment_name = 'batch_q'
            if self.config.psi_learning:
                experiment_name += '/psi_learning'
            if self.config.monte_carlo_count > 1:
                experiment_name = 'monte_carlo_targets/' + experiment_name

        # Folder for type of rewards used
        extra_save_dir = self.config.extra_save_dir
        if not extra_save_dir:
            if len(self.config.rewards) == 1:
                extra_save_dir = self.config.rewards[0]
            else:
                extra_save_dir = 'reward_combo'

        # Folder for which model was used
        extra_model_desc = ""
        if self.config.context_input_only:
            extra_model_desc = 'input_only_'
        if self.config.emotion and 'input_only' not in extra_model_desc:
            extra_model_desc += "emotion_"
        if self.config.infersent and 'input_only' not in extra_model_desc:
            extra_model_desc += "infersent_"

        # Make save path
        self.save_dir = default_save_path.joinpath(
            self.q_net.config.data, experiment_name, extra_save_dir,
            extra_model_desc + self.q_net.config.model, time_now)

        # Make directory and save config
        print("Saving output to", self.save_dir)
        os.makedirs(self.save_dir, exist_ok=True)
        with open(os.path.join(self.save_dir, 'config.txt'), 'w') as f:
            print(self.config, file=f)

        # Make loggers
        self.writer = TensorboardWriter(self.save_dir)
        self.pandas_path = os.path.join(self.save_dir, "metrics.csv")
        self.df = pd.DataFrame()

    def save_model(self, t):
        """Save parameters to checkpoint"""
        ckpt_path = os.path.join(self.save_dir, f'q_net{t}.pkl')
        print(f'Save parameters to {ckpt_path}')
        torch.save(self.q_net.model.state_dict(), ckpt_path)

        ckpt_path = os.path.join(self.save_dir, f'target_q_net{t}.pkl')
        torch.save(self.target_q_net.model.state_dict(), ckpt_path)

    def load_models(self):
        """Load parameters from RL checkpoint"""
        # Override specific checkpoint with particular one
        base_checkpoint_dir = str(Path(self.config.checkpoint).parent)
        if self.config.rl_ckpt_epoch is not None:
            q_ckpt_path = os.path.join(
                base_checkpoint_dir, 'q_net' + str(
                self.config.rl_ckpt_epoch) + '.pkl')
            target_q_ckpt_path = os.path.join(
                base_checkpoint_dir, 'target_q_net' + str(
                self.config.rl_ckpt_epoch) + '.pkl')
            self.t = int(self.config.rl_ckpt_epoch)
        else:
            ckpt_file = self.config.checkpoint.replace(base_checkpoint_dir, '')
            ckpt_file = ckpt_file.replace('/', '')
            ckpt_num = ckpt_file[len('q_net'):ckpt_file.find('.')]
            q_ckpt_path = self.config.checkpoint
            target_q_ckpt_path = os.path.join(
                base_checkpoint_dir, 'target_q_net' + ckpt_num + '.pkl')
            self.t = int(ckpt_num)

        print(f'Loading parameters for Q net from {q_ckpt_path}')
        q_ckpt = torch.load(q_ckpt_path)
        q_ckpt = convert_old_checkpoint_format(q_ckpt)
        self.q_net.model.load_state_dict(q_ckpt)

        print(f'Loading parameters for target Q net from {target_q_ckpt_path}')
        target_q_ckpt = torch.load(target_q_ckpt_path)
        target_q_ckpt = convert_old_checkpoint_format(target_q_ckpt)
        self.target_q_net.model.load_state_dict(target_q_ckpt)

        # Ensure weights are initialized to be on the GPU when necessary
        if torch.cuda.is_available():
            print('Converting checkpointed model to cuda tensors')
            self.q_net.model.cuda()
            self.target_q_net.model.cuda()

    def get_data_loader(self):
        # If checkpoint is for an emotion model, load that pickle file
        emotion_sentences = None
        if self.config.emotion:
            emotion_sentences = load_pickle(self.val_config.emojis_path)

        # Load infersent embeddings if necessary
        infersent_sentences = None
        if self.config.infersent:
            print('Loading infersent sentence embeddings...')
            infersent_sentences = load_pickle(self.val_config.infersent_path)
            embedding_size = infersent_sentences[0][0].shape[0]
            self.config.infersent_output_size = embedding_size
            self.val_config.infersent_output_size = embedding_size

        return get_loader(
            sentences=load_pickle(self.val_config.sentences_path),
            conversation_length=load_pickle(self.val_config.conversation_length_path),
            sentence_length=load_pickle(self.val_config.sentence_length_path),
            vocab=self.vocab,
            batch_size=self.config.batch_size,
            emojis=emotion_sentences,
            infersent=infersent_sentences)

    def interact(self):
        print("Commencing interaction with bot trained with RL")
        self.q_net.interact(
                max_sentence_length=self.config.max_sentence_length,
                max_conversation_length=self.config.max_conversation_length,
                sample_by='priority', debug=True, print_history=True)

def load_pickle(path):
    with open(path, 'rb') as f:
        return pickle.load(f)


import logging
_logger = logging.getLogger("theano.tensor.type")

import numpy

import theano
from theano import config
from theano.gof import Constant, hashtype, Type, Variable
from theano.gof.python25 import any
from theano.gof.utils import MethodNotDefined
from theano import scalar as scal

class TensorType(Type):
    """Symbolic `Type` representing a numpy.ndarray value."""

    filter_checks_isfinite = False
    """
    When this is True, strict filtering rejects data containing NaN or
    Inf entries. (Used in `DebugMode`)
    """

    def __init__(self, dtype, broadcastable, name=None, sparse_grad=False):
        """Initialize self.dtype and self.broadcastable.

        :Parameters:
         - `dtype`: str corresponding to numpy dtype (e.g., 'int64')
           The value (ndarray) associated to a `Variable` of this `Type` will
           have this dtype.
         - `broadcastable`: tuple, list, or array of boolean values
           This argument serves two purposes.  First, the True elements of this
           list indicate the dimensions where the shape of an associated value
           must be 1.  Secondly, the length of this list is the number of
           dimensions that an associated value must have.  See
           :doc:`broadcasting` for an explanation of how this list is used.
         - `name`: str
           Optional name for this type.
        """
        self.dtype = str(dtype)
        if self.dtype == 'floatX':
            self.dtype = config.floatX
        ###    broadcastable is immutable, and all elements are either
        ###    True or False
        self.broadcastable = tuple(bool(b) for b in broadcastable)
        self.dtype_specs()  # error checking is done there
        self.name = name
        self.numpy_dtype = numpy.dtype(self.dtype)
        self.sparse_grad = sparse_grad
        if sparse_grad:
            warnings.warn(
                "DEPRECATION WARNING: You use an old interface to"
                " AdvancedSubtensor1 sparse_grad. Now use"
                " theano.sparse_grad(a_tensor[an_int_vector]).")

    def filter(self, data, strict=False, allow_downcast=None):
        """Convert `data` to something which can be associated to a
        `TensorVariable`.

        This function is not meant to be called in user code.  It is for
        `Linker` instances to use when running a compiled graph.
        """
        # Explicit error message when one accidentally uses a Variable as
        # input (typical mistake, especially with shared variables).
        if isinstance(data, Variable):
            raise TypeError(
                    'Expected an array-like object, but found a Variable: '
                    'maybe you are trying to call a function on a (possibly '
                    'shared) variable instead of a numeric array?')

        if ((type(data) is numpy.ndarray)
                and (data.dtype == self.numpy_dtype)):
            if data.dtype.num != self.numpy_dtype.num:
                data = theano._asarray(data, dtype=self.dtype)
            # -- now fall through to ndim check
        elif((type(data) is numpy.memmap)
                and (data.dtype == self.numpy_dtype)):
            # numpy.memmap is a "safe" subclass of ndarray,
            # so we can use it whereever we expect a base ndarray.
            # however, casting it would defeat the purpose of not
            # loading the whole data into memory
            pass
        elif strict:
            # If any of the two conditions above was not met,
            # we raise a meaningful TypeError.
            if not (type(data) is numpy.ndarray):
                raise TypeError("%s expected a ndarray object." % self,
                        data, type(data))
            if data.dtype != self.numpy_dtype:
                raise TypeError(("%s expected a ndarray object with "
                        "dtype = %s (got %s).") % (
                            self, self.numpy_dtype, data.dtype))
            assert False, "This point should never be reached."
        else:
            if allow_downcast:
                # Convert to self.dtype, regardless of the type of data
                data = theano._asarray(data, dtype=self.dtype)
                # TODO: consider to pad shape with ones to make it consistent
                # with self.broadcastable... like vector->row type thing
            else:
                if isinstance(data, numpy.ndarray):
                    # Check if self.dtype can accurately represent data
                    # (do not try to convert the data)
                    up_dtype = scal.upcast(self.dtype, data.dtype)
                    if up_dtype == self.dtype:
                        # Bug in the following line when data is a
                        # scalar array, see
                        # http://projects.scipy.org/numpy/ticket/1611
                        # data = data.astype(self.dtype)
                        data = theano._asarray(data, dtype=self.dtype)
                    if up_dtype != self.dtype:
                        err_msg = (
                            '%s cannot store a value of dtype %s without '
                            'risking loss of precision. If you do not mind '
                            'this loss, you can: '
                            '1) explicitly cast your data to %s, or '
                            '2) set "allow_input_downcast=True" when calling '
                            '"function".'
                            % (self, data.dtype, self.dtype))
                        raise TypeError(err_msg, data)
                elif (allow_downcast is None and
                        type(data) is float and
                        self.dtype == theano.config.floatX):
                    # Special case where we allow downcasting of Python float
                    # literals to floatX, even when floatX=='float32'
                    data = theano._asarray(data, self.dtype)
                else:
                    # data has to be converted.
                    # Check that this conversion is lossless
                    converted_data = theano._asarray(data, self.dtype)
                    # We use the `values_eq` static function from TensorType
                    # to handle NaN values.
                    if TensorType.values_eq(numpy.asarray(data),
                                            converted_data,
                                            force_same_dtype=False):
                        data = converted_data
                    else:
                        # Do not print a too long description of data
                        # (ndarray truncates it, but it's not sure for data)
                        str_data = str(data)
                        if len(str_data) > 80:
                            str_data = str_data[:75] + '(...)'

                        err_msg = (
                            '%s cannot store accurately value %s, '
                            'it would be represented as %s. '
                            'If you do not mind this precision loss, you can: '
                            '1) explicitly convert your data to a numpy array '
                            'of dtype %s, or '
                            '2) set "allow_input_downcast=True" when calling '
                            '"function".'
                            % (self, data, converted_data, self.dtype))
                        raise TypeError(err_msg, data)

        if self.ndim != data.ndim:
            raise TypeError("Wrong number of dimensions: expected %s,"
                            " got %s with shape %s." % (self.ndim, data.ndim,
                                                        data.shape))
        if not data.flags.aligned:
            try:
                msg = "object buffer" + str(data.data)
            except AttributeError:
                msg = ""
            raise TypeError("The numpy.ndarray object is not aligned."
                            " Theano C code does not support that.",
                            msg,
                            "object shape", data.shape,
                            "object strides", data.strides,
                            "object dtype", data.dtype)

        i = 0
        for b in self.broadcastable:
            if b and data.shape[i] != 1:
                raise TypeError("Non-unit value on shape on a broadcastable"
                                " dimension.", data.shape, self.broadcastable)
            i += 1
        if (self.filter_checks_isfinite and
            not numpy.all(numpy.isfinite(data))):
            raise ValueError("non-finite elements not allowed")
        return data

    def filter_variable(self, other):
        """Convert a symbolic Variable into a TensorType, if compatible.

        For the moment, only a TensorType or CudaNdarrayType will be
        converted, provided they have the same number of dimensions,
        broadcastable pattern, and dtype.
        """
        if hasattr(other, '_as_TensorVariable'):
            other = other._as_TensorVariable()

        if not isinstance(other, Variable):
            # The value is not a Variable: we cast it into
            # a Constant of the appropriate Type.
            other = self.Constant(type=self, data=other)

        if other.type == self:
            return other

        raise TypeError(
                'Cannot convert Type %(othertype)s '
                '(of Variable %(other)s) into Type %(self)s. '
                'You can try to manually convert %(other)s into a %(self)s.'
                % dict(
                    othertype=other.type,
                    other=other,
                    self=self)
                )

    def value_validity_msg(self, a):
        try:
            self.filter(a, strict=True)
        except Exception, e:
            return str(e)
        return "value is valid"

    def dtype_specs(self):
        """Return a tuple (python type, c type, numpy typenum) that corresponds
        to self.dtype.

        This function is used internally as part of C code generation.
        """
        # TODO: add more type correspondances for e.g. int32, int64, float32,
        # complex64, etc.
        try:
            return {
                'float32': (float, 'npy_float32', 'NPY_FLOAT32'),
                'float64': (float, 'npy_float64', 'NPY_FLOAT64'),
                'uint8': (int, 'npy_uint8', 'NPY_UINT8'),
                'int8': (int, 'npy_int8', 'NPY_INT8'),
                'uint16': (int, 'npy_uint16', 'NPY_UINT16'),
                'int16': (int, 'npy_int16', 'NPY_INT16'),
                'uint32': (int, 'npy_uint32', 'NPY_UINT32'),
                'int32': (int, 'npy_int32', 'NPY_INT32'),
                'uint64': (int, 'npy_uint64', 'NPY_UINT64'),
                'int64': (int, 'npy_int64', 'NPY_INT64'),
                'complex128': (complex, 'theano_complex128', 'NPY_COMPLEX128'),
                'complex64': (complex, 'theano_complex64', 'NPY_COMPLEX64')
                }[self.dtype]
        except KeyError:
            raise TypeError("Unsupported dtype for %s: %s"
                    % (self.__class__.__name__, self.dtype))

    def to_scalar_type(self):
        return scal.get_scalar_type(dtype=self.dtype)

    def __eq__(self, other):
        """Compare True iff other is the same kind of TensorType"""
        return type(self) == type(other) and other.dtype == self.dtype \
            and other.broadcastable == self.broadcastable

    @staticmethod
    def may_share_memory(a, b):
        # This is a method of TensorType, so both a and b should be ndarrays
        if isinstance(a, numpy.ndarray) and isinstance(b, numpy.ndarray):
            return numpy.may_share_memory(a, b)
        else:
            return False

    @staticmethod
    def values_eq(a, b, force_same_dtype=True):
        # TODO: check to see if the shapes must match
        #      for now, we err on safe side...
        if a.shape != b.shape:
            return False
        if force_same_dtype and a.dtype != b.dtype:
            return False
        a_eq_b = (a == b)
        r = numpy.all(a_eq_b)
        if r:
            return True
        # maybe the trouble is that there are NaNs
        a_missing = numpy.isnan(a)
        if a_missing.any():
            b_missing = numpy.isnan(b)
            return numpy.all(a_eq_b + (a_missing == b_missing))
        else:
            return False

    @staticmethod
    def values_eq_approx(a, b, allow_remove_inf=False, allow_remove_nan=False,
                         rtol=None, atol=None):
        """
        :param allow_remove_inf: If True, when there is an inf in a,
                                 we allow any value in b in that position.
                                 Event -inf
        :param allow_remove_nan: If True, when there is a nan in a,
                                 we allow any value in b in that position.
                                 Event +-inf
        :param rtol: relative tolerance, passed to _allclose
        :param atol: absolute tolerance, passed to _allclose
        """
        if isinstance(a, numpy.ndarray) and isinstance(b, numpy.ndarray):
            if a.shape != b.shape:
                return False
            if a.dtype != b.dtype:
                return False
            if 'int' in str(a.dtype):
                return numpy.all(a == b)
            else:
                # work around a numpy.allclose bug:
                # http://projects.scipy.org/numpy/ticket/1672
                if a.ndim == 0 and numpy.isinf(a):
                    a = a.reshape(1)
                    b = b.reshape(1)

                cmp = theano.tensor.basic._allclose(a, b, rtol=rtol, atol=atol)
                if cmp:
                    # Numpy claims they are close, this is good enough for us.
                    return True
                # Numpy is unhappy, but it does not necessarily mean that a and
                # b are different. Indeed, Numpy does not like missing values
                # and will return False whenever some are found in a or b.
                # The proper way would be to use the MaskArray stuff available
                # in Numpy. However, it looks like it has been added to Numpy's
                # core recently, so it may not be available to everyone. Thus,
                # for now we use a home-made recipe, that should probably be
                # revisited in the future.
                a_missing = numpy.isnan(a)
                a_inf = numpy.isinf(a)

                if not (a_missing.any() or (allow_remove_inf and a_inf.any())):
                    # There are no missing values in a, thus this is not the
                    # reason why numpy.allclose(a, b) returned False.
                    _logger.info(
                        'numpy allclose failed for abs_err %f and rel_err %f',
                        numpy.max(abs(a - b)),
                        numpy.max(abs(a - b) / (abs(a) + abs(b))))
                    return False
                # The following line is what numpy.allclose bases its decision
                # upon, according to its documentation.
                rtol = 1.0000000000000001e-05
                atol = 1e-8
                cmp_elemwise = (numpy.absolute(a - b) <=
                        (atol + rtol * numpy.absolute(b)))
                # Find places where both a and b have missing values.
                both_missing = a_missing * numpy.isnan(b)

                # Find places where both a and b have inf of the same sign.
                both_inf = a_inf * numpy.isinf(b)

                # cmp_elemwise is weird when we have inf and -inf.
                # set it to False
                cmp_elemwise = numpy.where(
                        both_inf & cmp_elemwise,
                        a == b,
                        cmp_elemwise)

                # check the sign of the inf
                both_inf = numpy.where(both_inf, (a == b), both_inf)

                if allow_remove_inf:
                    both_inf += a_inf
                if allow_remove_nan:
                    both_missing += a_missing

                # Combine all information.
                return (cmp_elemwise + both_missing + both_inf).all()

        return False

    @staticmethod
    def values_eq_approx_remove_inf(a, b):
        return TensorType.values_eq_approx(a, b, True)

    @staticmethod
    def values_eq_approx_remove_nan(a, b):
        return TensorType.values_eq_approx(a, b, False, True)

    @staticmethod
    def values_eq_approx_remove_inf_nan(a, b):
        return TensorType.values_eq_approx(a, b, True, True)

    def __hash__(self):
        """Hash equal for same kinds of TensorType"""
        return hashtype(self) ^ hash(self.dtype) ^ hash(self.broadcastable)

    ndim = property(lambda self: len(self.broadcastable),
            doc="number of dimensions")
    """Number of dimensions

    This read-only property is the preferred way to get the number of
    dimensions of a `TensorType`.

    """

    def make_variable(self, name=None):
        """Return a `TensorVariable` of this type

        :Parameters:
         - `name`: str
           A pretty name to identify this `Variable` when printing and
           debugging
        """
        return self.Variable(self, name=name)

    def __str__(self):
        if self.name:
            return self.name
        else:
            b = self.broadcastable
            named_broadcastable = {(): 'scalar',
                     (False,): 'vector',
                     (False, True): 'col',
                     (True, False): 'row',
                     (False, False): 'matrix'}
            if b in named_broadcastable:
                bcast = named_broadcastable[b]
            else:
                if any(b):
                    bcast = str(b)
                else:
                    bcast = '%iD' % len(b)
            return "TensorType(%s, %s)" % (str(self.dtype), bcast)

    def __repr__(self):
        return str(self)
        #"TensorType{%s, %s}" % (str(self.dtype), str(self.broadcastable))

    def c_declare(self, name, sub, check_input=True):
        """Override `CLinkerType.c_declare` """
        if(check_input):
            check = """
            typedef %(dtype)s dtype_%(name)s;
            """ % dict(sub, name=name, dtype=self.dtype_specs()[1])
        else:
            check = ""
        declaration = """
        PyArrayObject* %(name)s;
        """ % dict(sub, name=name, dtype=self.dtype_specs()[1])

        return declaration + check

    def c_init(self, name, sub):
        """Override `CLinkerType.c_init` """
        return """
        %(name)s = NULL;
        """ % dict(sub, name=name, type_num=self.dtype_specs()[2])

    def c_extract(self, name, sub, check_input=True):
        """Override `CLinkerType.c_extract` """
        if(check_input):
            check = """
            %(name)s = NULL;
            if (py_%(name)s == Py_None) {
                // We can either fail here or set %(name)s to NULL and rely on Ops
                // using tensors to handle the NULL case, but if they fail to do so
                // they'll end up with nasty segfaults, so this is public service.
                PyErr_SetString(PyExc_ValueError, "expected an ndarray, not None");
                %(fail)s
            }
            if (!PyArray_Check(py_%(name)s)) {
                PyErr_SetString(PyExc_ValueError, "expected an ndarray");
                %(fail)s
            }
            // We expect %(type_num)s
            if (!PyArray_ISALIGNED((PyArrayObject*) py_%(name)s)) {
                PyArrayObject * tmp = (PyArrayObject*) py_%(name)s;
                PyErr_Format(PyExc_NotImplementedError,
                             "expected an aligned array of type %%ld "
                             "(%(type_num)s), got non-aligned array of type %%ld"
                             " with %%ld dimensions, with 3 last dims "
                             "%%ld, %%ld, %%ld"
                             " and 3 last strides %%ld %%ld, %%ld.",
                             (long int) %(type_num)s,
                             (long int) PyArray_TYPE((PyArrayObject*) py_%(name)s),
                             (long int) PyArray_NDIM(tmp),
                             (long int) PyArray_NDIM(tmp) >= 3 ?
            PyArray_DIMS(tmp)[PyArray_NDIM(tmp)-3] : -1,
                             (long int) PyArray_NDIM(tmp) >= 2 ?
            PyArray_DIMS(tmp)[PyArray_NDIM(tmp)-2] : -1,
                             (long int) PyArray_NDIM(tmp) >= 1 ?
            PyArray_DIMS(tmp)[PyArray_NDIM(tmp)-1] : -1,
                             (long int) PyArray_NDIM(tmp) >= 3 ?
            PyArray_STRIDES(tmp)[PyArray_NDIM(tmp)-3] : -1,
                             (long int) PyArray_NDIM(tmp) >= 2 ?
            PyArray_STRIDES(tmp)[PyArray_NDIM(tmp)-2] : -1,
                             (long int) PyArray_NDIM(tmp) >= 1 ?
            PyArray_STRIDES(tmp)[PyArray_NDIM(tmp)-1] : -1
            );
                %(fail)s
            }
            // This is a TypeError to be consistent with DEBUG_MODE
            // Note: DEBUG_MODE also tells the name of the container
            if (PyArray_TYPE((PyArrayObject*) py_%(name)s) != %(type_num)s) {
                PyErr_Format(PyExc_TypeError,
                             "expected type_num %%d (%(type_num)s) got %%d",
                             %(type_num)s, PyArray_TYPE((PyArrayObject*) py_%(name)s));
                %(fail)s
            }
            """ % dict(sub, name=name, type_num=self.dtype_specs()[2])
        else:
            check = ""
        return check + """
        %(name)s = (PyArrayObject*)(py_%(name)s);
        Py_XINCREF(%(name)s);
        """ % dict(sub, name=name, type_num=self.dtype_specs()[2])

    def c_cleanup(self, name, sub):
        """Override `CLinkerType.c_cleanup` """
        return """
        if (%(name)s) {
            Py_XDECREF(%(name)s);
        }
        """ % locals()

    def c_sync(self, name, sub):
        """Override `CLinkerType.c_sync` """
        fail = sub['fail']
        type_num = self.dtype_specs()[2]
        return """
        {Py_XDECREF(py_%(name)s);}
        if (!%(name)s) {
            Py_INCREF(Py_None);
            py_%(name)s = Py_None;
        }
        else if ((void*)py_%(name)s != (void*)%(name)s) {
            py_%(name)s = (PyObject*)%(name)s;
        }

        {Py_XINCREF(py_%(name)s);}

        if (%(name)s && !PyArray_ISALIGNED((PyArrayObject*) py_%(name)s)) {
            PyErr_Format(PyExc_NotImplementedError,
                         "c_sync: expected an aligned array, got non-aligned array of type %%ld"
                         " with %%ld dimensions, with 3 last dims "
                         "%%ld, %%ld, %%ld"
                         " and 3 last strides %%ld %%ld, %%ld.",
                         (long int) PyArray_TYPE((PyArrayObject*) py_%(name)s),
                         (long int) PyArray_NDIM(%(name)s),
                         (long int) PyArray_NDIM(%(name)s) >= 3 ?
        PyArray_DIMS(%(name)s)[PyArray_NDIM(%(name)s)-3] : -1,
                         (long int) PyArray_NDIM(%(name)s) >= 2 ?
        PyArray_DIMS(%(name)s)[PyArray_NDIM(%(name)s)-2] : -1,
                         (long int) PyArray_NDIM(%(name)s) >= 1 ?
        PyArray_DIMS(%(name)s)[PyArray_NDIM(%(name)s)-1] : -1,
                         (long int) PyArray_NDIM(%(name)s) >= 3 ?
        PyArray_STRIDES(%(name)s)[PyArray_NDIM(%(name)s)-3] : -1,
                         (long int) PyArray_NDIM(%(name)s) >= 2 ?
        PyArray_STRIDES(%(name)s)[PyArray_NDIM(%(name)s)-2] : -1,
                         (long int) PyArray_NDIM(%(name)s) >= 1 ?
        PyArray_STRIDES(%(name)s)[PyArray_NDIM(%(name)s)-1] : -1
        );
            %(fail)s
        }
        """ % locals()

    def c_headers(self):
        """Override `CLinkerObject.c_headers` """
        return scal.get_scalar_type(self.dtype).c_headers()

    def c_libraries(self):
        return scal.get_scalar_type(self.dtype).c_libraries()

    def c_compile_args(self):
        return scal.get_scalar_type(self.dtype).c_compile_args()

    def c_support_code(self):
        """Override `CLinkerObject.c_support_code` """
        return scal.get_scalar_type(self.dtype).c_support_code()

    def c_init_code(self):
        return scal.get_scalar_type(self.dtype).c_init_code()

    def c_code_cache_version(self):
        scalar_version = scal.get_scalar_type(self.dtype).c_code_cache_version()
        if scalar_version:
            return (11,) + scalar_version
        else:
            return ()

    def value_zeros(self, shape):
        """
        Create an numpy ndarray full of 0 values.
        """
        return numpy.zeros(shape, dtype=self.dtype)

    def get_shape_info(self, obj):
        """
        Return the information needed to compute the memory size of ``obj``.

        The memory size is only the data, so this excludes the container.
        For an ndarray, this is the data, but not the ndarray object and
        other data structures such as shape and strides.

        ``get_shape_info()`` and ``get_size()`` work in tandem for the memory
        profiler.

        ``get_shape_info()`` is called during the execution of the function.
        So it is better that it is not too slow.

        ``get_size()`` will be called on the output of this function
        when printing the memory profile.

        :param obj: The object that this Type represents during execution
        :return: Python object that ``self.get_size()`` understands
        """
        return obj.shape

    def get_size(self, shape_info):
        """ Number of bytes taken by the object represented by shape_info.

        :param shape_info: the output of the call to get_shape_info()
        :return: the number of bytes taken by the object described by
            ``shape_info``.
        """
        if shape_info:
            return numpy.prod(shape_info) * numpy.dtype(self.dtype).itemsize
        else:  # a scalar
            return numpy.dtype(self.dtype).itemsize
theano.compile.ops.expandable_types += (TensorType,)

# Register TensorType C code for ViewOp.
theano.compile.register_view_op_c_code(
        TensorType,
        """
        Py_XDECREF(%(oname)s);
        %(oname)s = %(iname)s;
        Py_XINCREF(%(oname)s);
        """,
        version=1)

# Register TensorType C code for Shape Op.
theano.compile.register_shape_c_code(
    TensorType,
    """
    npy_intp shape[] = {PyArray_NDIM(%(iname)s)};
    if(%(oname)s == NULL || (PyArray_DIMS(%(oname)s)[0] != shape[0]))
    {
        Py_XDECREF(%(oname)s);
        %(oname)s = (PyArrayObject*) PyArray_SimpleNew(1, shape, NPY_INT64);
    }
    for(int i=0;i<shape[0];i++)
    {
        ((npy_int64*)PyArray_GETPTR1(%(oname)s, i))[0] = PyArray_DIMS(%(iname)s)[i];
    }
    """,
    version=1)

# Register TensorType C code for ViewOp.
theano.compile.register_shape_i_c_code(
        TensorType,
        """
        if(!%(oname)s)
            %(oname)s=(PyArrayObject*)PyArray_EMPTY(0, NULL, NPY_INT64, 0);
        ((npy_int64*)PyArray_DATA(%(oname)s))[0]=PyArray_DIMS(%(iname)s)[%(i)s];
        """,
        """
        if (%(i)s>=PyArray_NDIM(%(iname)s)){
            PyErr_SetString(PyExc_TypeError,
                "Number of dimensions lower than expected");
            %(fail)s
        }
        """,
        version=3)

# Register TensorType C code for DeepCopyOp
theano.compile.register_deep_copy_op_c_code(
        TensorType,
        """
        int alloc = %(oname)s == NULL;
        for(int i=0; !alloc && i<PyArray_NDIM(%(oname)s); i++) {
           if(PyArray_DIMS(%(iname)s)[i] != PyArray_DIMS(%(oname)s)[i]) {
               alloc = true;
               break;
           }
        }
        if(alloc) {
            Py_XDECREF(%(oname)s);
            %(oname)s = (PyArrayObject*)PyArray_NewCopy(%(iname)s,
                                                        NPY_ANYORDER);
            if (!%(oname)s)
            {
                PyErr_SetString(PyExc_ValueError,
                                "DeepCopyOp: the copy failed!");
                %(fail)s;
            }
        } else {
            if(PyArray_CopyInto(%(oname)s, %(iname)s)){
                PyErr_SetString(PyExc_ValueError,
            "DeepCopyOp: the copy failed into already allocated space!");
                %(fail)s;
            }
        }
        """,
        version=2)

theano.compile.register_rebroadcast_c_code(
    TensorType,
    """
    if(PyArray_DIMS(%(iname)s)[%(axis)s] != 1){
        PyErr_Format(PyExc_ValueError,
            "Dimension %(axis)s in Rebroadcast's input was"
            " supposed to be 1 (got %%d instead)",
            PyArray_DIMS(%(iname)s)[%(axis)s]);
        %(fail)s
    }
    """,
        version=1)

theano.compile.register_specify_shape_c_code(
    TensorType,
    """
        if (PyArray_NDIM(%(iname)s) != PyArray_DIMS(%(shape)s)[0]) {
            PyErr_Format(PyExc_AssertionError,
                         "SpecifyShape: vector of shape has %%d elements,"
                         " but the input has %%d dimensions.",
                         PyArray_NDIM(%(iname)s),
                         PyArray_DIMS(%(shape)s)[0]);
            %(fail)s;
        }
        for(int i = 0; i < PyArray_NDIM(%(iname)s); i++){
            dtype_%(shape)s shp = ((dtype_%(shape)s*)PyArray_GETPTR1(%(shape)s,
                                                                     i))[0];
            if (PyArray_DIMS(%(iname)s)[i] != shp) {
                PyErr_Format(PyExc_AssertionError,
                             "SpecifyShape: dim %%d of input has shape %%d,"
                             " expected %%d.",
                             i, PyArray_DIMS(%(iname)s)[i],
                             shp);
                %(fail)s;
            }
        }
        Py_XDECREF(%(oname)s);
        %(oname)s = %(iname)s;
        Py_XINCREF(%(oname)s);
    """,
    version=1)


# -*- coding: utf-8 -*-

import numpy as np

# make colormaps
#from matplotlib.colors import LinearSegmentedColormap
import matplotlib.pyplot as plt
esrBlueRed = {'red':   ((0.0, 0.0, 0.0),
                        (0.5, 0.0, 0.0),
                        (1.0, 1.0, 1.0)),

              'green': ((0.0, 0.0, 0.0),
                        (0.5, 0.0, 0.0),
                        (1.0, 0.0, 0.0)),

              'blue':  ((0.0, 1.0, 1.0),
                        (0.5, 0.0, 0.0),
                        (1.0, 0.0, 0.0))
              }

isomorphicTest = {'red':   ((0.0, 0.0, 0.0),
                            (0.33, 0.0, 0.0),
                            (0.67, 1.0, 1.0),
                            (1, 1.0, 1.0)),

              'green': ((0.0, 0.0, 0.0),
                            (0.33, 0.0, 0.0),
                            (0.67, 0.0, 0.0),
                            (1, 1.0, 1.0)),

              'blue':  ((0.0, 0.0, 0.0),
                            (0.33, 1.0, 1.0),
                            (0.67, 0.0, 0.0),
                            (1, 1.0, 1.0)),
              }

esrJet = {'red':       ((0./12, 0.0, 0.0),
                        (1./12, 0.0, 0.0),
                        (2./12, 0.0, 0.0),
                        (3./12, 0.0, 0.0),
                        (4./12, 0.0, 0.0),
                        (4.7/12, 0.5, 0.5),
                        (6./12, 1.0, 1.0),
                        (7.3/12, 1.0, 1.0),
                        (8./12, 1.0, 1.0),
                        (9./12, 1.0, 1.0),
                        (10./12, 1.0, 1.0),
                        (11./12, 1.0, 1.0),
                        (12./12, 1.0, 1.0)),

          'green':     ((0./12, 0.0, 0.0),
                        (1./12, 0.0, 0.0),
                        (2./12, 0.0, 0.0),
                        (3./12, 0.5, 0.5),
                        (4.7/12, 1.0, 1.0),
                        (5./12, 1.0, 1.0),
                        (6./12, 1.0, 1.0),
                        (7.3/12, 0.5, 0.5),
                        (8./12, 0.0, 0.0),
                        (9./12, 0.0, 0.0),
                        (10./12, 0.0, 0.0),
                        (11./12, 0.5, 0.5),
                        (12./12, 1.0, 1.0)),

          'blue':      ((0./12, 0.0, 0.0),
                        (1./12, 0.5, 0.5),
                        (2./12, 1.0, 1.0),
                        (3./12, 0.5, 0.5),
                        (4.7/12, 0.0, 0.0),
                        (5./12, 0.0, 0.0),
                        (6./12, 0.0, 0.0),
                        (7.3/12, 0.0, 0.0),
                        (8./12, 0.0, 0.0),
                        (9./12, 0.5, 0.5),
                        (10./12, 1.0, 1.0),
                        (11./12, 1.0, 1.0),
                        (12./12, 1.0, 1.0)),
              }

_nipy_spectral_pinktop = {
      'red': [(0.0, 0.0, 0.0), (0.05, 0.4667, 0.4667),
              (0.10, 0.5333, 0.5333), (0.15, 0.0, 0.0),
              (0.20, 0.0, 0.0), (0.25, 0.0, 0.0),
              (0.30, 0.0, 0.0), (0.35, 0.0, 0.0),
              (0.40, 0.0, 0.0), (0.45, 0.0, 0.0),
              (0.50, 0.0, 0.0), (0.55, 0.0, 0.0),
              (0.60, 0.0, 0.0), (0.65, 0.7333, 0.7333),
              (0.70, 0.9333, 0.9333), (0.75, 1.0, 1.0),
              (0.80, 1.0, 1.0), (0.85, 1.0, 1.0),
              (0.90, 0.8667, 0.8667), (0.95, 0.80, 0.80),
              (1.0, 1.0, 1.0)],
    'green': [(0.0, 0.0, 0.0), (0.05, 0.0, 0.0),
              (0.10, 0.0, 0.0), (0.15, 0.0, 0.0),
              (0.20, 0.0, 0.0), (0.25, 0.4667, 0.4667),
              (0.30, 0.6000, 0.6000), (0.35, 0.6667, 0.6667),
              (0.40, 0.6667, 0.6667), (0.45, 0.6000, 0.6000),
              (0.50, 0.7333, 0.7333), (0.55, 0.8667, 0.8667),
              (0.60, 1.0, 1.0), (0.65, 1.0, 1.0),
              (0.70, 0.9333, 0.9333), (0.75, 0.8000, 0.8000),
              (0.80, 0.6000, 0.6000), (0.85, 0.0, 0.0),
              (0.90, 0.0, 0.0), (0.95, 0.0, 0.0),
              (1.0, 0.5, 0.5)],
     'blue': [(0.0, 0.0, 0.0), (0.05, 0.5333, 0.5333),
              (0.10, 0.6000, 0.6000), (0.15, 0.6667, 0.6667),
              (0.20, 0.8667, 0.8667), (0.25, 0.8667, 0.8667),
              (0.30, 0.8667, 0.8667), (0.35, 0.6667, 0.6667),
              (0.40, 0.5333, 0.5333), (0.45, 0.0, 0.0),
              (0.5, 0.0, 0.0), (0.55, 0.0, 0.0),
              (0.60, 0.0, 0.0), (0.65, 0.0, 0.0),
              (0.70, 0.0, 0.0), (0.75, 0.0, 0.0),
              (0.80, 0.0, 0.0), (0.85, 0.0, 0.0),
              (0.90, 0.0, 0.0), (0.95, 0.0, 0.0),
              (1.0, 1.0, 1.0)],
}

# no pink at top
esrJet2 = {'red':       ((0, 0.0, 0.0),  # black
                        (0.25, 0.0, 0.0),  # blue
                        (0.5, 0.0, 0.0),  # green
                        (0.75, 1.0, 1.0),  # yellow
                        (1.0, 1.0, 1.0)),  # red

          'green':     ((0, 0.0, 0.0),  # black
                        (0.25, 0.0, 0.0),  # blue
                        (0.5, 1.0, 1.0),  # green
                        (0.75, 1.0, 1.0),  # yellow
                        (1.0, 0.0, 0.0)),  # red

          'blue':      ((0, 0.0, 0.0),  # black
                        (0.25, 1.0, 1.0),  # blue
                        (0.5, 0.0, 0.0),  # green
                        (0.75, 0.0, 0.0),  # yellow
                        (1.0, 0.0, 0.0)),  # red
              }

test =  {'red':      ((0.0, 0.0, 0.0),
                        (0.5, 0.5, 0.5),
                        (1.0, 1.0, 1.0)),

          'green':     ((0.0, 0.0, 0.0),
                        (0.5, 0.25, 0.25),
                        (1.0, 0.0, 0.0)),

          'blue':      ((0.0, 0.0, 0.0),
                        (0.5, 0.25, 0.25),
                        (1.0, 0.0, 0.0)),
              }

plt.register_cmap(name='isomorphicTest', data=isomorphicTest)

#blue_red1 = LinearSegmentedColormap('BlueRed1', cdict1)
plt.register_cmap(name='esrBlueRed', data=esrBlueRed)
plt.register_cmap(name='esrJet', data=esrJet)
plt.register_cmap(name='nipy_spectral_pinktop', data=_nipy_spectral_pinktop)
plt.register_cmap(name='esrJet2', data=esrJet2)

#@profile
def load_param_single_simple(fn, status=[0, np.inf], trueAzEl=False):
    """
    Loads parameters from a single GUISDAP result file. NOT meant to be a
    complete replacement for GUISDAPs own load_param.m. Loads physical
    parameters and time, az/el etc. and little else from monostatic experiments.

    Parameters
    ----------

    fn : string, required
        name of file to load

    status : list of length 2
        [max_status, max_residual]
        Status: 0 = OK, 1 = max number of iterations exceeded, 2 = No fit done
        because data too noisy

    trueAzEl : boolean
        If False, azimuth and elevation will be cast into 0-360 and 0-90 degrees

    """

    import datetime as dt
    from scipy.io import loadmat

    params = loadmat(fn, mat_dtype=True)

    # needed for calculations and other computations
    Te_Ti = params['r_param'][:, 2]  # Te/Ti
    errTe_Ti = params['r_error'][:, 2]
    r_status = params['r_status'][:, 0]

    # don't really know what this is
    if 'r_Offsetppd' in params:
        rOff = params['r_Offsetppd']
    elif 'r_phasepush' in params:
        rOff = params['r_phasepush']
    else:
        rOff = np.nan

    c1 = r_status > status[0]
    c2 = params['r_res'][:, 0] > status[1]
    c3 = params['r_error'][:, :8] > 0
    c12 = np.array([np.bitwise_or(c1, c2), ]*8).T

    params['r_error'][:, :8][c12*c3] = np.nan
    params['r_param'][c12*c3] = np.nan

    # Time
    # XXX: Time is a python datetime object, not MATLAB datenum!
    t = params['r_time']
    tStart = dt.datetime(int(t[0, 0]), int(t[0, 1]), int(t[0, 2]), int(t[0, 3]), int(t[0, 4]), int(t[0, 5]), int(round(np.mod(t[0, 5], 1)*1e6, 0)))
    tEnd = dt.datetime(int(t[1, 0]), int(t[1, 1]), int(t[1, 2]), int(t[1, 3]), int(t[1, 4]), int(t[1, 5]), int(round(np.mod(t[1, 5], 1)*1e6, 0)))

    # 1D params
    Az = params['r_az'][0][0]  # azimuth
    El = params['r_el'][0][0]  # elevation
    Pt = params['r_Pt'][0][0]/10000  # transmitter power
    Tsys = np.median(params['r_Tsys'])
    Oppd_Php = rOff

    # 2D params
    Ran = params['r_range'][:, 0]  # range of scattering volume
    Alt = params['r_h'][:, 0]  # altitude of scattering volume
    Ne = params['r_param'][:, 0]
    Ti = params['r_param'][:, 1]
    Te = Te_Ti * Ti
    Vi = -params['r_param'][:, 4]
    Coll = params['r_param'][:, 3]  # collision frequency
    Comp = params['r_dp'][:, 0]  # Ion composition ([O+]/Ne)
    Res = params['r_res'][:, 0]  # residual of the fit (or standard deviation?)

    # 2D errors
    errNe = params['r_error'][:, 0]
    errTi = params['r_error'][:, 1]
    errTe = (errTi/Ti + errTe_Ti/Te_Ti)*Te
    errVi = params['r_error'][:, 4]
    errColl = params['r_error'][:, 3]

    Time = np.array([[tStart, tEnd], ]).T
    par1D = np.array([[Az, El, Pt, Tsys], ])
    if not np.isnan(Oppd_Php):
        par1D = np.append(par1D, Oppd_Php, axis=1)
    par2D = np.column_stack((Ran, Alt, Ne, Te, Ti, Vi, Coll, Comp, Res))
    par2D = np.expand_dims(par2D, 1)
    err2D = np.column_stack((errNe, errTe, errTi, errVi, errColl))
    err2D = np.expand_dims(err2D, 1)
    rpar2D = np.array([])  # XXX currently not implemented

    # cast azimuth and elevation into range 0-360, 0-90 degrees
    if not trueAzEl:
        d = np.where(par1D[:, 1] > 90)
        par1D[d, 1] = 180 - par1D[d, 1]
        par1D[d, 0] = par1D[d, 0] + 180
        par1D[:, 0] = np.mod(par1D[:, 0]+360, 360)

    return Time, par2D, par1D, rpar2D, err2D

def load_param_simple(path, trueAzEl=False):
    """
    Loads parameters from a directory of GUISDAP result files. NOT meant to be
    a complete replacement for GUISDAPs own load_param.m. Loads physical
    parameters and time, az/el etc. and little else from monostatic experiments.

    Parameters
    ----------

    path: string, required
        name of file to load

    trueAzEl : boolean
        If False, azimuth and elevation will be cast into 0-360 and 0-90 degrees

    Returns
    -------

    * `N` is number of integrations (i.e., number of files read)
    * `M` is number of range gates

    Time : 2D array (2, `N`)
        First index is start timestamp (0) and end timestamp (1) of integration.
    par2D : 3D array (`M`, `N`, 9)
        Third index is parameter: Range (0), Altitude (1), Ne (2), Te (3),
        Ti (4), Vi (5), Coll (6), Comp (7), Res (8)
    par1D : 2D array (5, `N`)
        First index is parameter: Az (0), El (1), Pt (2), Tsys (3), Oppd/Php (4)
    rpar2D : Empty array (not implemented)
        ..
    err2D : 3D array (`M`, `N`, 5)
        Errors for the following parameters (third index): Ne (0), Te (1), Ti (2), Vi (3), Coll (4)

    """

    import os
    import fnmatch

    mat_files = fnmatch.filter(os.listdir(path), '*.mat')
    n_ip = len(mat_files)  # number of integration periods

    try:
        import frogress
        iterator = frogress.bar(enumerate(mat_files), steps=len(mat_files))
    except:
        iterator = enumerate(mat_files)

    for i, mat_file in iterator:
        s_Time, s_par2D, s_par1D, s_rpar2D, s_err2D = load_param_single_simple(os.path.join(path, mat_file), trueAzEl=trueAzEl)

        if i == 0:  # initialize data structures
            n_ran = len(s_par2D[:, 0, 1])  # number of range gates

            Time = np.empty((2, n_ip), dtype=object)
            par2D = np.empty((n_ran, n_ip, 9))*np.nan
            par1D = np.zeros((n_ip, s_par1D.shape[1]))
            rpar2D = np.array([])
            err2D = np.empty((n_ran, n_ip, 5))*np.nan

        # correct dimensions if number of range gates have changed
        # XXX: Double-check how this is done in the original matlab code
        if par2D[:, 0, 0].shape[0] < s_par2D[:, 0, 0].shape[0]:
            par2D = np.append(par2D, np.ones((s_par2D.shape[0]-par2D.shape[0], par2D.shape[1], par2D.shape[2]))*np.nan, axis=0)
            err2D = np.append(err2D, np.ones((s_err2D.shape[0]-err2D.shape[0], err2D.shape[1], err2D.shape[2]))*np.nan, axis=0)
        elif par2D[:, 0, 0].shape[0] > s_par2D[:, 0, 0].shape[0]:
            s_par2D = np.append(s_par2D, np.ones((par2D.shape[0]-s_par2D.shape[0], s_par2D.shape[1], s_par2D.shape[2]))*np.nan, axis=0)
            s_err2D = np.append(s_err2D, np.ones((err2D.shape[0]-s_err2D.shape[0], s_err2D.shape[1], s_err2D.shape[2]))*np.nan, axis=0)

        # somewhat the same for par1D
        if s_par1D.shape[1] > par1D.shape[1]:
            par1D = np.vstack((par1D, np.empty((par1D.shape[1], 1))*np.nan))
        elif s_par1D.shape[1] < par1D.shape[1]:
            s_par1D = np.append(s_par1D, np.nan)

        # add current data to data structures
        Time[:, i] = s_Time[:, 0]
        par1D[i, :] = s_par1D[0, :]
        par2D[:, i, :] = s_par2D[:, 0, :]
        err2D[:, i, :] = s_err2D[:, 0, :]

    return Time, par2D, par1D, rpar2D, err2D

def gg2gc(gg):
    """transforms coordinates from geographic (lat, lon, h) to geocentric"""

    import math

    factor = math.pi/180  # conversion factor from degrees to radians
    r_earth = 6378.135  # earth radius (km) and flatness factor
    g = 1.00673944  # earth flatness factor

    lat = gg[0]*factor
    lon = gg[1]*factor
    h = gg[2]

    hor = (r_earth/math.sqrt(1+math.tan(lat)**2/g)+h*math.cos(lat))
    gc = [hor*math.cos(lon), hor*math.sin(lon), r_earth/math.sqrt(g+g**2/math.tan(lat)**2)+h*math.sin(lat)]

    return gc

def gc2gg(gc):
    """transforms coordinates from geocentric to geographic (lat, lon, h)"""

    import math

    factor = math.pi/180  # conversion factor from degrees to radians
    r_earth = 6378.135  # earth radius (km) and flatness factor
    g = 1.00673944  # earth flatness factor

    gg = [0, 0, 0]  # initialize gg, not needed in original MATLAB code...

    if gc[0] == 0 and gc[1] == 0:
        print('Beware of the spinning earth axis!')
        gg = [90, 0, gc[2] - r_earth/g]
    else:
        gg[1] = math.atan2(gc[1], gc[0]) / factor
        r0 = math.sqrt(sum([gc[0]*gc[0], gc[1]*gc[1]]))
        xi0 = gc[2] / (r0 * math.sqrt(g))
        xi_iter = r_earth*(g-1)/(g*r0)
        tanxi = xi0
        tanxi = xi0 + xi_iter * tanxi / math.sqrt(1+tanxi**2)
        tanxi = xi0 + xi_iter * tanxi / math.sqrt(1+tanxi**2)
        gg[0] = math.atan(math.sqrt(g)*tanxi)/factor
        gg[2] = math.sqrt(1+g*tanxi**2)*(r0-r_earth/math.sqrt(1+tanxi**2))

    return gg

def loc2gg(site1, loc):
    """transforms the scattering point location given in local coordinates
    loc   [elevation, azimuth, range] at location
    site1 [latitude, longitude, height]  to geographic coordinates
    """

    import math
    import numpy as np

    factor = math.pi/180

    #  first calculate the  transformation matrices
    lat1 = site1[0]*factor
    lon1 = site1[1]*factor
    sinlat = math.sin(lat1)
    coslat = math.cos(lat1)
    sinlon = math.sin(lon1)
    coslon = math.cos(lon1)
    rlocgc = np.array([[sinlat*coslon, -sinlon,  coslat*coslon],
                       [sinlat*sinlon,  coslon,  coslat*sinlon],
                       [-coslat,             0,  sinlat]])

    s1 = loc[0]*factor
    s2 = loc[1]*factor
    s3 = loc[2]
    loc = s3*np.array([-math.cos(s2)*math.cos(s1), math.sin(s2)*math.cos(s1), math.sin(s1)]).T
    gc_site1 = gg2gc(site1)  # Site1 to geogentric
    gc_sp = gc_site1 + np.dot(rlocgc, loc).T  # Add scattering distance in geocentric
    gg_sp = gc2gg(gc_sp)  # Transform back to geographic

    return gg_sp