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

'''Non-relativistic UKS analytical nuclear gradients'''

import time
import numpy
import scipy.linalg
from pyscf import lib
from pyscf.lib import logger
from pyscf.grad import rhf as rhf_grad
from pyscf.grad import rks as rks_grad
from pyscf.grad import uhf as uhf_grad
from pyscf.dft import numint, gen_grid
from pyscf import __config__

def get_veff(ks_grad, mol=None, dm=None):
    '''Coulomb + XC functional
    '''
    if mol is None: mol = ks_grad.mol
    if dm is None: dm = ks_grad.base.make_rdm1()
    t0 = (time.clock(), time.time())

    mf = ks_grad.base
    ni = mf._numint
    if ks_grad.grids is not None:
        grids = ks_grad.grids
    else:
        grids = mf.grids
    if grids.coords is None:
        grids.build(with_non0tab=True)

    if mf.nlc != '':
        raise NotImplementedError
    #enabling range-separated hybrids
    omega, alpha, hyb = ni.rsh_and_hybrid_coeff(mf.xc, spin=mol.spin)

    mem_now = lib.current_memory()[0]
    max_memory = max(2000, ks_grad.max_memory*.9-mem_now)
    if ks_grad.grid_response:
        exc, vxc = get_vxc_full_response(ni, mol, grids, mf.xc, dm,
                                         max_memory=max_memory,
                                         verbose=ks_grad.verbose)
        logger.debug1(ks_grad, 'sum(grids response) %s', exc.sum(axis=0))
    else:
        exc, vxc = get_vxc(ni, mol, grids, mf.xc, dm,
                           max_memory=max_memory, verbose=ks_grad.verbose)
    t0 = logger.timer(ks_grad, 'vxc', *t0)

    if abs(hyb) < 1e-10:
        vj = ks_grad.get_j(mol, dm)
        vxc += vj[0] + vj[1]
    else:
        vj, vk = ks_grad.get_jk(mol, dm)
        vk *= hyb
        if abs(omega) > 1e-10:  # For range separated Coulomb operator
            with mol.with_range_coulomb(omega):
                vk += ks_grad.get_k(mol, dm) * (alpha - hyb)
        vxc += vj[0] + vj[1] - vk

    return lib.tag_array(vxc, exc1_grid=exc)

def get_vxc(ni, mol, grids, xc_code, dms, relativity=0, hermi=1,
            max_memory=2000, verbose=None):
    xctype = ni._xc_type(xc_code)
    make_rho, nset, nao = ni._gen_rho_evaluator(mol, dms, hermi)
    ao_loc = mol.ao_loc_nr()

    vmat = numpy.zeros((2,3,nao,nao))
    if xctype == 'LDA':
        ao_deriv = 1
        for ao, mask, weight, coords \
                in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
            rho_a = make_rho(0, ao[0], mask, 'LDA')
            rho_b = make_rho(1, ao[0], mask, 'LDA')
            vxc = ni.eval_xc(xc_code, (rho_a,rho_b), 1, relativity, 1, verbose)[1]
            vrho = vxc[0]
            aow = numpy.einsum('pi,p->pi', ao[0], weight*vrho[:,0])
            rks_grad._d1_dot_(vmat[0], mol, ao[1:4], aow, mask, ao_loc, True)
            aow = numpy.einsum('pi,p->pi', ao[0], weight*vrho[:,1])
            rks_grad._d1_dot_(vmat[1], mol, ao[1:4], aow, mask, ao_loc, True)
            rho = vxc = vrho = aow = None

    elif xctype == 'GGA':
        ao_deriv = 2
        for ao, mask, weight, coords \
                in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
            rho_a = make_rho(0, ao[:4], mask, 'GGA')
            rho_b = make_rho(1, ao[:4], mask, 'GGA')
            vxc = ni.eval_xc(xc_code, (rho_a,rho_b), 1, relativity, 1, verbose)[1]
            wva, wvb = numint._uks_gga_wv0((rho_a,rho_b), vxc, weight)

            rks_grad._gga_grad_sum_(vmat[0], mol, ao, wva, mask, ao_loc)
            rks_grad._gga_grad_sum_(vmat[1], mol, ao, wvb, mask, ao_loc)
            rho_a = rho_b = vxc = wva = wvb = None

    elif xctype == 'NLC':
        raise NotImplementedError('NLC')
    else:
        raise NotImplementedError('meta-GGA')

    exc = numpy.zeros((mol.natm,3))
    # - sign because nabla_X = -nabla_x
    return exc, -vmat

def get_vxc_full_response(ni, mol, grids, xc_code, dms, relativity=0, hermi=1,
                          max_memory=2000, verbose=None):
    '''Full response including the response of the grids'''
    xctype = ni._xc_type(xc_code)
    make_rho, nset, nao = ni._gen_rho_evaluator(mol, dms, hermi)
    ao_loc = mol.ao_loc_nr()
    aoslices = mol.aoslice_by_atom()

    excsum = 0
    vmat = numpy.zeros((2,3,nao,nao))
    if xctype == 'LDA':
        ao_deriv = 1
        for atm_id, (coords, weight, weight1) \
                in enumerate(rks_grad.grids_response_cc(grids)):
            ngrids = weight.size
            sh0, sh1 = aoslices[atm_id][:2]
            mask = gen_grid.make_mask(mol, coords)
            ao = ni.eval_ao(mol, coords, deriv=ao_deriv, non0tab=mask)
            rho_a = make_rho(0, ao[0], mask, 'LDA')
            rho_b = make_rho(1, ao[0], mask, 'LDA')
            exc, vxc = ni.eval_xc(xc_code, (rho_a,rho_b), 1, relativity, 1, verbose)[:2]
            vrho = vxc[0]

            vtmp = numpy.zeros((3,nao,nao))
            aow = numpy.einsum('pi,p->pi', ao[0], weight*vrho[:,0])
            rks_grad._d1_dot_(vtmp, mol, ao[1:4], aow, mask, ao_loc, True)
            vmat[0] += vtmp
            excsum += numpy.einsum('r,r,nxr->nx', exc, rho_a+rho_b, weight1)
            excsum[atm_id] += numpy.einsum('xij,ji->x', vtmp, dms[0]) * 2

            vtmp = numpy.zeros((3,nao,nao))
            aow = numpy.einsum('pi,p->pi', ao[0], weight*vrho[:,1])
            rks_grad._d1_dot_(vtmp, mol, ao[1:4], aow, mask, ao_loc, True)
            vmat[1] += vtmp
            excsum[atm_id] += numpy.einsum('xij,ji->x', vtmp, dms[1]) * 2
            rho = vxc = vrho = aow = None

    elif xctype == 'GGA':
        ao_deriv = 2
        for atm_id, (coords, weight, weight1) \
                in enumerate(rks_grad.grids_response_cc(grids)):
            ngrids = weight.size
            sh0, sh1 = aoslices[atm_id][:2]
            mask = gen_grid.make_mask(mol, coords)
            ao = ni.eval_ao(mol, coords, deriv=ao_deriv, non0tab=mask)
            rho_a = make_rho(0, ao[:4], mask, 'GGA')
            rho_b = make_rho(1, ao[:4], mask, 'GGA')
            exc, vxc = ni.eval_xc(xc_code, (rho_a,rho_b), 1, relativity, 1, verbose)[:2]
            wva, wvb = numint._uks_gga_wv0((rho_a,rho_b), vxc, weight)

            vtmp = numpy.zeros((3,nao,nao))
            rks_grad._gga_grad_sum_(vtmp, mol, ao, wva, mask, ao_loc)
            vmat[0] += vtmp
            excsum += numpy.einsum('r,r,nxr->nx', exc, rho_a[0]+rho_b[0], weight1)
            excsum[atm_id] += numpy.einsum('xij,ji->x', vtmp, dms[0]) * 2

            vtmp = numpy.zeros((3,nao,nao))
            rks_grad._gga_grad_sum_(vtmp, mol, ao, wvb, mask, ao_loc)
            vmat[1] += vtmp
            excsum[atm_id] += numpy.einsum('xij,ji->x', vtmp, dms[1]) * 2
            rho_a = rho_b = vxc = wva = wvb = None

    elif xctype == 'NLC':
        raise NotImplementedError('NLC')
    else:
        raise NotImplementedError('meta-GGA')

    # - sign because nabla_X = -nabla_x
    return excsum, -vmat

class Gradients(uhf_grad.Gradients):

    grid_response = getattr(__config__, 'grad_uks_Gradients_grid_response', False)

    def __init__(self, mf):
        uhf_grad.Gradients.__init__(self, mf)
        self.grids = None
        self.grid_response = False
        self._keys = self._keys.union(['grid_response', 'grids'])

    def dump_flags(self, verbose=None):
        uhf_grad.Gradients.dump_flags(self, verbose)
        logger.info(self, 'grid_response = %s', self.grid_response)
        return self

    get_veff = get_veff

    def extra_force(self, atom_id, envs):
        '''Hook for extra contributions in analytical gradients.

        Contributions like the response of auxiliary basis in density fitting
        method, the grid response in DFT numerical integration can be put in
        this function.
        '''
        if self.grid_response:
            vhf = envs['vhf']
            log = envs['log']
            log.debug('grids response for atom %d %s',
                      atom_id, vhf.exc1_grid[atom_id])
            return vhf.exc1_grid[atom_id]
        else:
            return 0

Grad = Gradients

from pyscf import dft
dft.uks.UKS.Gradients = dft.uks_symm.UKS.Gradients = lib.class_as_method(Gradients)

if __name__ == '__main__':
    from pyscf import gto
    from pyscf import dft

    mol = gto.Mole()
    mol.atom = [
        ['O' , (0. , 0.     , 0.)],
        [1   , (0. , -0.757 , 0.587)],
        [1   , (0. ,  0.757 , 0.587)] ]
    mol.basis = '631g'
    mol.charge = 1
    mol.spin = 1
    mol.build()
    mf = dft.UKS(mol)
    mf.conv_tol = 1e-12
    #mf.grids.atom_grid = (20,86)
    e0 = mf.scf()
    g = mf.Gradients()
    print(lib.finger(g.kernel()) - -0.12090786243525126)
#[[-5.23195019e-16 -5.70291415e-16  5.32918387e-02]
# [ 1.33417513e-16  6.75277008e-02 -2.66519852e-02]
# [ 1.72274651e-16 -6.75277008e-02 -2.66519852e-02]]
    g.grid_response = True
    print(lib.finger(g.kernel()) - -0.12091122429043633)
#[[-2.95956939e-16 -4.22275612e-16  5.32998759e-02]
# [ 1.34532051e-16  6.75279140e-02 -2.66499379e-02]
# [ 1.68146089e-16 -6.75279140e-02 -2.66499379e-02]]

    mf.xc = 'b88,p86'
    e0 = mf.scf()
    g = Gradients(mf)
    print(lib.finger(g.kernel()) - -0.11509739136150157)
#[[ 2.58483362e-16  5.82369026e-16  5.17616036e-02]
# [-5.46977470e-17  6.39273304e-02 -2.58849008e-02]
# [ 5.58302713e-17 -6.39273304e-02 -2.58849008e-02]]
    g.grid_response = True
    print(lib.finger(g.kernel()) - -0.11507986316077731)

    mf.xc = 'b3lypg'
    e0 = mf.scf()
    g = Gradients(mf)
    print(lib.finger(g.kernel()) - -0.10202554999695367)
#[[ 6.47874920e-16 -2.75292214e-16  3.97215970e-02]
# [-6.60278148e-17  5.87909340e-02 -1.98650384e-02]
# [ 6.75500259e-18 -5.87909340e-02 -1.98650384e-02]]

    mol = gto.Mole()
    mol.atom = [
        ['H' , (0. , 0. , 1.804)],
        ['F' , (0. , 0. , 0.   )], ]
    mol.unit = 'B'
    mol.basis = '631g'
    mol.charge = -1
    mol.spin = 1
    mol.build()

    mf = dft.UKS(mol)
    mf.conv_tol = 1e-14
    mf.kernel()
    print(lib.finger(Gradients(mf).kernel()) - 0.10365160440876001)
# sum over z direction non-zero, due to meshgrid response
# H    -0.0000000000     0.0000000000    -0.1481125370
# F    -0.0000000000     0.0000000000     0.1481164667
    mf = dft.UKS(mol)
    mf.grids.prune = None
    mf.grids.level = 6
    mf.conv_tol = 1e-14
    mf.kernel()
    print(lib.finger(Gradients(mf).kernel()) - 0.10365040148752827)
# H     0.0000000000     0.0000000000    -0.1481124925
# F    -0.0000000000     0.0000000000     0.1481122913


# This source reads a text file and produces a data set
# The file is of the form (where [denotes optional and repeatable]:
# X,Y,Z[,ARRAYNAME]
# x0,y0,z0[,value]
# x1,y1,z1[,value
# ...
# from paraview import vtk
# from random import randint as randi
# pdo=self.GetOutputDataObject(0)
# pdi=self.GetInputDataObject(0,0)
# pdo.CopyAttributes(pdi)
# pdo.Allocate(500,1)
# maxPt=pdi.GetNumberOfPoints()
# for i in range(500):
# 	aLine = vtk.vtkLine()
# 	aLine.GetPointIds().SetId(0, randi(0, maxPt))
# 	aLine.GetPointIds().SetId(1, randi(0, maxPt))
# 	pdo.InsertNextCell(aLine.GetCellType(), aLine.GetPointIds())
from paraview import vtk
from paraview.vtk import dataset_adapter as DA
import os, sys
from numpy import array, single, zeros
import numpy
import os, sys

def loadFoam(cObj, cNum):
    cFile = "/Users/maderk/Documents/MATLAB/TextureTensor/Foam/ujax/dk31m00/"
    cFile += "bubble_dk31_%02d.csv" % cNum
    filterList = []
    # filterList+=[('MASK_DISTANCE_MEAN',lambda x: x>180)]
    # filterList+=[('VOLUME',lambda x: x<4000)]
    synList = []
    synList += [
        (
            "Alignment",
            lambda x: 180
            / numpy.pi
            * calcAlignment(x, "PCA1_X", "PCA1_Y", "PCA1_Z", "VOLUME"),
        )
    ]
    synList += [
        ("ShapeAnisotropy", lambda x: calcAnisotropy(x, "PCA1_S", "PCA2_S", "PCA3_S"))
    ]
    synList += [
        ("ShapeOblateness", lambda x: calcOblateness(x, "PCA1_S", "PCA2_S", "PCA3_S"))
    ]
    synList += [("PosTheta", lambda x: calcTheta(x, "POS_X", "POS_Y", recenter=True))]
    synList += [("Theta", lambda x: calcTheta(x, "PCA1_X", "PCA2_Y", recenter=False))]
    synList += [
        ("Theta-XZ", lambda x: chkgvlcalcTheta(x, "PCA1_X", "PCA2_Z", recenter=False))
    ]
    return ImportCSV(cObj, cFile, filterList, synList)

def initCol(cName, data, keepPoints):
    name = vtk.vtkTypeFloat32Array()
    name.SetNumberOfComponents(1)
    name.SetNumberOfTuples(0)
    name.SetName(cName)
    if keepPoints is None:
        keepPoints = range(0, len(data))
    for cI in keepPoints:
        name.InsertNextTuple((data[cI],))
    return name

def initCol3(cName, dat1, dat2, dat3, keepPoints=None):
    name = vtk.vtkTypeFloat32Array()
    name.SetNumberOfComponents(3)
    name.SetNumberOfTuples(0)
    name.SetName(cName)
    if keepPoints is None:
        keepPoints = range(0, len(dat1))
    for cI in keepPoints:
        name.InsertNextTuple((dat1[cI], dat2[cI], dat3[cI]))
    print (cName, len(keepPoints))
    return name

def initColTens(
    cName, dat1, dat2, dat3, dat4, dat5, dat6, dat7, dat8, dat9, keepPoints=None
):
    name = vtk.vtkTypeFloat32Array()
    name.SetNumberOfComponents(9)
    name.SetNumberOfTuples(0)
    name.SetName(cName + "_T")
    if keepPoints is None:
        keepPoints = range(0, len(dat1))

    for cI in keepPoints:
        name.InsertNextTuple(
            (
                dat1[cI],
                dat2[cI],
                dat3[cI],
                dat4[cI],
                dat5[cI],
                dat6[cI],
                dat7[cI],
                dat8[cI],
                dat9[cI],
            )
        )
    print (cName, len(keepPoints))
    return name

def getTriplets(allCols, cn1="_X", cn2="_Y", cn3="_Z"):
    xcols = filter(lambda x: (x.find(cn1) >= 0), allCols)
    xcols = filter(lambda x: not (x.find("T" + cn1) >= 0), xcols)
    repPrefix = lambda x, y, z: z.join(x.split(y))
    goodList = [
        (
            repPrefix(curCol, cn1, ""),
            (
                repPrefix(curCol, cn1, cn1),
                repPrefix(curCol, cn1, cn2),
                repPrefix(curCol, cn1, cn3),
            ),
        )
        for curCol in xcols
        if (
            (repPrefix(curCol, cn1, cn2) in allCols)
            and (repPrefix(curCol, cn1, cn3) in allCols)
        )
    ]
    return goodList

def getTensors(
    allCols, cn1=["_XX", "_XY", "_XZ", "_YX", "_YY", "_YZ", "_ZX", "_ZY", "_ZZ"]
):
    xcols = filter(lambda x: x.find("T" + cn1[0]) >= 0, allCols)
    repPrefix = lambda x, y, z: z.join(x.split(y))
    goodList = []
    for curCol in xcols:
        cRow = []
        for curEle in cn1:
            cName = repPrefix(curCol, cn1[0], curEle)
            if cName in allCols:
                cRow += [cName]
            else:
                cRow += ["NULL"]
        if len(cRow) > 2:
            goodList += [(repPrefix(curCol, cn1[0], ""), tuple(cRow))]
    return goodList

def wtexture(xv, yv, zv, wv):
    """ Since the function will be used alot and is tedious to implement """
    xv = numpy.array(xv)
    yv = numpy.array(yv)
    zv = numpy.array(zv)
    wv = numpy.array(wv)

    swv = sum(wv)
    xm = numpy.multiply(xv, wv) / swv
    ym = numpy.multiply(yv, wv) / swv
    zm = numpy.multiply(zv, wv) / swv

    xv -= xm
    yv -= ym
    zv -= zm

    xv *= wv
    yv *= wv
    zv *= wv

    textMat = [[sum(xv * xv) / swv, sum(xv * yv) / swv, sum(xv * zv) / swv]]
    textMat += [[sum(yv * xv) / swv, sum(yv * yv) / swv, sum(yv * zv) / swv]]
    textMat += [[sum(zv * xv) / swv, sum(zv * yv) / swv, sum(zv * zv) / swv]]
    return textMat
    DA.numpyTovtkDataArray(numpy.random.rand(50, 1))

import numpy as np

class ctexture:
    def __init__(self, textMat):
        self.textMat = textMat
        (self._w, self._v) = np.linalg.eigh(textMat)
        # self._w=np.sqrt(self._w) # No sqrt in texture tensor

    def todict(self, name="TEXTURE", length=1, oDict=None):
        if oDict is None:
            oDict = {}
        for (i, cAxI) in enumerate("XYZ"):
            for (j, cAxJ) in enumerate("XYZ"):
                oDict[name + "T_" + cAxI + cAxJ] = [self.textMat[i, j]] * length
        return oDict

    def mean(self):
        return sqrt(self.textMat.trace())

    def v1(self):
        return self._v[:, 2]

    def alignment(self):
        return 3.0 / 2.0 * (self._w[2] / sum(self._w) - 1.0 / 3.0) * 100

class lacDB_TEXT:
    def __init__(self):
        self.count = 0.0
        self.countn = 0
        self.sumx2 = 0.0
        self.sumx = 0.0
        self.sumxy = 0.0
        self.sumxz = 0.0
        self.sumy2 = 0.0
        self.sumyz = 0.0
        self.sumz2 = 0.0
        self.text = 0

    def step(self, valueX, valueY, valueZ, valueW=1.0):
        if valueX is None:
            return 0
        if valueY is None:
            return 0
        if valueZ is None:
            return 0
        if valueW is None:
            return 0
        self.count += valueW
        self.countn += 1
        self.sumx += valueW * valueX
        self.sumy += valueW * valueY
        self.sumz += valueW * valueZ

        self.sumx2 += valueW * valueX ** 2
        self.sumxy += valueW * valueX * valueY
        self.sumxz += valueW * valueX * valueZ
        self.sumy2 += valueW * valueY ** 2
        self.sumyz += valueW * valueY * valueZ
        self.sumz2 += valueW * valueZ ** 2

    def textMat(self):
        meanx = self.sumx / self.count
        meany = self.sumy / self.count
        meanz = self.sumz / self.count

        return np.array(
            [
                [
                    self.sumx2 / self.count - meanx ** 2,
                    self.sumxy / self.count - meanx * meany,
                    self.sumxz / self.count - meanx * meanz,
                ],
                [
                    self.sumxy / self.count - meanx * meany,
                    self.sumy2 / self.count - meany ** 2,
                    self.sumyz / self.count - meany * meanz,
                ],
                [
                    self.sumxz / self.count - meany * meanz,
                    self.sumyz / self.count - meany * meanz,
                    self.sumz2 / self.count - meanz ** 2,
                ],
            ]
        )

    def finalize(self):
        if self.countn < 2:
            return 100
        self.text = ctexture(self.textMat())
        return self.text.alignment()

def calcAlignment(
    allcols, cn1, cn2, cn3, cnWeight=None, returnAxis=True, signedAxis=True
):
    """ Since the function will be used alot and is tedious to implement """
    ptlist = []
    ptlist += [allcols[cn1]]
    ptlist += [allcols[cn2]]
    ptlist += [allcols[cn3]]
    ptlist = numpy.array(ptlist)
    ptmag = numpy.sqrt(ptlist[0, :] ** 2 + ptlist[1, :] ** 2 + ptlist[2, :] ** 2)
    ptlist[0, :] /= ptmag
    ptlist[1, :] /= ptmag
    ptlist[2, :] /= ptmag
    if cnWeight is None:
        textMat = numpy.cov(ptlist)
    else:
        xv = ptlist[0, :]
        yv = ptlist[1, :]
        zv = ptlist[2, :]
        wv = numpy.array(allcols[cnWeight])
        textMat = wtexture(xv, yv, zv, wv)
        print (len(xv), len(wv))
        print textMat
    (_w, _v) = numpy.linalg.eigh(textMat)
    print "Main Direction:" + str(_v[:, 2]) + ", Overall alignment:" + str(
        3.0 / 2.0 * (_w[2] / sum(_w) - 1.0 / 3.0) * 100
    )

    if returnAxis:
        oMat = []
        sMat = []
        for cRow in range(3):
            mv = numpy.array([list(_v[:, cRow])] * ptlist.shape[1]).transpose()
            oMat += [numpy.abs(numpy.sum(numpy.multiply(ptlist, mv), 0))]
            sMat += [numpy.sign(numpy.sum(numpy.multiply(ptlist, mv), 0))]
        oAx = numpy.argmax(numpy.array(oMat), 0)
        if signedAxis:
            sMat = numpy.array(sMat).transpose()
            return numpy.array(
                map(lambda x: 2 * (x[1] + 1) + 1 * (x[0][x[1]] == 1), zip(sMat, oAx))
            )
        else:
            return oAx + 1
    else:
        mv = numpy.array([list(_v[:, 2])] * ptlist.shape[1]).transpose()
        return numpy.arccos(numpy.abs(numpy.sum(numpy.multiply(ptlist, mv), 0)))

def calcAnisotropy(allcols, cn1, cn2, cn3):
    """ Since the function will be used alot and is tedious to implement """
    p1list = numpy.array(allcols[cn1])
    p2list = numpy.array(allcols[cn2])
    p3list = numpy.array(allcols[cn3])
    return 3.0 / 2.0 * (p1list / (p1list + p2list + p3list) - 1.0 / 3.0) * 100

def calcOblateness(allcols, cn1, cn2, cn3):
    """ Since the function will be used alot and is tedious to implement """
    p1list = numpy.array(allcols[cn1])
    p2list = numpy.array(allcols[cn2])
    p3list = numpy.array(allcols[cn3])
    return 2 * (p1list - p2list) / (p1list - p3list) - 1

def calcTheta(allcols, cnx, cny, cnw=None, recenter=True):
    """ Since the function will be used alot and is tedious to implement """
    p1list = numpy.array(allcols[cnx])
    p2list = numpy.array(allcols[cny])
    if cnw is not None:
        p3list = numpy.array(allcols[cnw])
    else:
        p3list = p1list * 0 + 1
    mx = numpy.multiply(p1list, p3list) / sum(p3list)
    my = numpy.multiply(p2list, p3list) / sum(p3list)
    if not recenter:
        mx = 0
        my = 0
    return 180 / numpy.pi * numpy.arctan2(p2list - my, p1list - mx)

def ImportCSV(cObj, filename="", filterList=[], synList=[], outObj=0, fromR=False):
    cTable = cObj.GetOutputDataObject(outObj)
    (header, cols) = LoadCSVFile(filename, fromR=fromR)
    return ImportData(cTable, cols, filterList, synList)

def ImportDB(cObj, sqlRest="", filterList=[], synList=[], outObj=0):
    cTable = cObj.GetOutputDataObject(outObj)
    if len(sqlRest) > 0:
        sqlRest = "WHERE " + sqlRest
    import dbImport as dbi

    cur = dbi.cur
    headerCols = [
        cCol[0] for cCol in cur.execute("SHOW COLUMNS FROM LACUNA").fetchall()
    ]
    cols = {}
    for cCol in headerCols:
        cols[cCol] = []
    for cResult in cur.execute(
        "SELECT " + ",".join(headerCols) + " FROM LACUNA " + sqlRest
    ):
        for (ccCol, cVal) in zip(headerCols, cResult):
            cols[ccCol] += [cVal]
    return ImportData(cTable, cols, filterList, synList)

def ImportData(output, cols, filterList=[], synList=[]):
    """ Filter list is for adding lambda based filters, synList is for synthetic columns """

    print "Importing Next Step..."
    # ncols=lacpa_adddb(None,cols,header,'',False) # formatted and scaled
    # cols=ncols
    print "Filtering List"
    keepPoints = numpy.array([True] * len(cols.values()[0]))  # only
    for (cFiltCol, cFiltFunc) in filterList:
        keepPoints &= map(cFiltFunc, cols[cFiltCol])
    print "Keeping - " + str((sum(keepPoints), len(keepPoints)))
    keepPoints = [cValue for (cValue, isKept) in enumerate(keepPoints) if isKept]
    for cKey in cols.keys():
        cols[cKey] = numpy.array(cols[cKey])[keepPoints]
    print " Adding columns"
    for (cName, cData) in cols.items():
        output.AddColumn(initCol(cName, cData, None))
    print " Adding Triplets"
    allTrips = getTriplets(cols.keys())
    print allTrips
    for (cName, cCols) in allTrips:
        output.AddColumn(
            initCol3(cName, cols[cCols[0]], cols[cCols[1]], cols[cCols[2]], None)
        )
    print " Adding Tensors"
    cols["NULL"] = 0 * single(cols.values()[0])
    allTens = getTensors(cols.keys())
    print allTens
    for (cName, cCols) in allTens:
        output.AddColumn(
            initColTens(
                cName,
                cols[cCols[0]],
                cols[cCols[1]],
                cols[cCols[2]],
                cols[cCols[3]],
                cols[cCols[4]],
                cols[cCols[5]],
                cols[cCols[6]],
                cols[cCols[7]],
                cols[cCols[8]],
                None,
            )
        )

    print " Adding synthetic columns"
    for (cName, cFunc) in synList:
        output.AddColumn(initCol(cName, cFunc(cols), None))
    return output

def addRowsToTable(output, junkName, outMat, entry_order):
    # globals()['a']=entry_order
    # globals()['b']=outMat
    outArray = numpy.array(outMat)
    outDict = {}
    for (cI, cLabel) in enumerate(entry_order):
        outDict[cLabel] = outArray[:, cI]
    return outDict

from numpy import single

def parseCSV(text, filename="", fromR=False):
    def temp_clean(text):
        return ("".join(text.split("/"))).upper().strip()

    def temp_parse(temp):
        ntemp = []
        errCount = 0
        for val in temp:
            if val.strip().upper() == "NAN":
                cval = "0"
                errCount += 1
            else:
                cval = val
            try:
                cval = single(cval)
            except:
                cval = -1
                errCount += 1
            ntemp += [cval]
        return (ntemp, errCount)

    rows = text.split("\n")
    # First row is header
    if fromR:
        fileDict = {}
        hRow = 0
    else:
        hRow = 1
        head1 = rows[0]
        head1 = "".join(head1.split("//"))
        newStr = [cEle.strip().split(":") for cEle in head1.strip().split(",")]
        fileDict = {}
        for cEle in newStr:
            if len(cEle) == 2:
                fileDict[temp_clean(cEle[0])] = cEle[1].split("/")[-1].strip()
    fTime = True
    head2 = rows[hRow]
    head2 = "".join(head2.split("//"))
    head2 = [temp_clean(cEle) for cEle in head2.strip().split(",")]
    # Check for duplicates in header string (and just use the last entry)
    # Generate a dictionary of all header entries
    cleanHeader = {}
    for k in range(len(head2)):
        cleanHeader[head2[k]] = k
    # create a new null filled header
    head2 = ["NULL"] * len(head2)
    # use the dictionary to repopulate the head2 entry
    for cKey in cleanHeader.keys():
        head2[cleanHeader[cKey]] = cKey
    outTable = {}
    for col in head2:
        outTable[col] = []
    for row in rows[2:]:
        temp = row.split(",")
        try:
            (ntemp, errs) = temp_parse(temp)
            if errs < 2:
                if len(ntemp) == len(head2):
                    for k in range(0, len(head2)):
                        outTable[head2[k]] += [ntemp[k]]
        except:
            # if fTime: print (len(ntemp),len(head2))
            fTime = False
            temp = []
    for col in head2:
        outTable[col] = numpy.array(outTable[col])
    outrows = len(outTable[head2[0]])
    print "Parsed .. " + str(outrows) + " of " + str(len(rows))
    return (outrows, fileDict, outTable)

def LoadCSVFile(filename, fromR=False):
    if True:
        rawtext = "".join(open(filename).readlines())
    else:
        print filename + " is garbage"
    if True:
        (outrows, a, b) = parseCSV(rawtext, fromR=fromR)
        if outrows > 2:
            return (a, b)
        else:
            print filename + " is too short!" + str((outrows, a, b))
    else:
        print filename + " is junk:" + rawtext[0:100]

def getProjNum(*args):
    return 1

def getSampleNum(*args):
    return 1

def lacpa_adddb(
    cur,
    ptList,
    oTable,
    rawFilename,
    processName=True,
    tableName="Lacuna",
    CanalMode=0,
    projectTitle="None",
):
    lacNumOffset = 0
    if processName:
        (filename, lacNumOffset) = processInputName(rawFilename, lacFilename)
    else:
        filename = rawFilename
    nptList = dict([(cCol.upper(), cDat) for (cCol, cDat) in ptList.items()])
    # ptList=CaseFreeDict(ptList)

    dbLen = len(ptList["SCALE_X"])
    if not oTable.has_key("SAMPLE"):
        oTable["SAMPLE"] = ""
        print filename + " is missing sample name"
    dx = numpy.median(ptList["SCALE_X"])
    dy = numpy.median(ptList["SCALE_Y"])
    dz = numpy.median(ptList["SCALE_Z"])
    dr = numpy.sqrt(dx ** 2 + dy ** 2 + dz ** 2)
    lacTemp = {}
    if type(projectTitle) is type(""):
        cProjNum = getProjNum(cur, projectTitle)
    else:
        cProjNum = projectTitle
    cSampleNum = getSampleNum(cur, filename, cProjNum)
    lacTemp["SAMPLE_AIM_Number"] = (cSampleNum,) * dbLen
    lacTemp["PROJECT"] = (cProjNum,) * dbLen
    lacunIds = [lacId + lacNumOffset for lacId in ptList["LACUNA_NUMBER"]]
    lacTemp[tableName + "_NUMBER"] = tuple(lacunIds)
    # Variables that scale directly with x,y,z voxel size
    lacTemp["VOX_SIZE"] = tuple(numpy.abs(ptList["SCALE_X"] * 1000))
    scaleVars = ["POS", "STD", "PROJ"]
    for cVar in scaleVars:
        for cAx in ["X", "Y", "Z"]:
            lacTemp[cVar + "_" + cAx] = tuple(
                ptList[cVar + "_" + cAx] * ptList["SCALE_" + cAx]
            )
    # This doesnt work since I dont save PCA1,2,3 dumb
    # Variables that scale with PCA 1,2,3 voxel size * denotes PCA1, PCA2, PCA3
    pcaScaleVars = ["*_S", "PROJ_*"]
    for cAx in ["PCA1", "PCA2", "PCA3"]:
        cDr = numpy.sqrt(
            (ptList[cAx + "_X"] * dx) ** 2
            + (ptList[cAx + "_Y"] * dy) ** 2
            + (ptList[cAx + "_Z"] * dz) ** 2
        )
        for cVar in pcaScaleVars:
            rcVar = cAx.join(cVar.split("*"))
            lacTemp[rcVar] = tuple(ptList[rcVar] * cDr)
    # Normal Variables
    normalVars = ["PCA1_X", "PCA1_Y", "PCA1_Z", "PCA2_X", "PCA2_Y", "PCA2_Z"]
    normalVars += ["MASK_GRAD_X", "MASK_GRAD_Y", "MASK_GRAD_Z", "MASK_ANGLE"]
    if CanalMode == 0:
        normalVars += ["Canal_ANGLE", "Canal_GRAD_X", "Canal_GRAD_Y", "Canal_GRAD_Z"]
    for cVar in normalVars:
        if ptList.has_key(cVar):
            lacTemp[cVar] = tuple(ptList[cVar])
        elif (cVar.find("GRAD") >= 0) | (cVar.find("ANGLE")):
            lacTemp[cVar] = (-1,) * dbLen
        else:
            print "Missing important column:" + cVar + ", what the frick!"
    # Variables that require a radial scaling factor
    radialVars = ["MASK_DISTANCE_MEAN", "MASK_DISTANCE_STD"]  # 'MASK_DISTANCE_COV'
    radialVars += ["OBJ_RADIUS", "OBJ_RADIUS_STD"]
    if CanalMode == 0:
        if ptList.has_key(cVar):
            radialVars += ["Canal_DISTANCE_MEAN", "Canal_DISTANCE_STD"]
    for cVar in radialVars:
        if ptList.has_key(cVar):
            lacTemp[cVar] = tuple(numpy.abs(ptList[cVar] * dr))
    # Variables that require a radial cubed scaling factor
    volVars = ["VOLUME", "VOLUME_BOX"]
    for cVar in volVars:
        lacTemp[cVar] = tuple(numpy.abs(ptList[cVar] * dx * dy * dz))
    if ptList.has_key("SHELL_CNT"):
        lacTemp["VOLUME_LAYER"] = tuple(
            numpy.abs((ptList["VOLUME"] - ptList["SHELL_CNT"]) * dx * dy * dz)
        )
    # GrayAnalysis Columns
    if ptList.has_key("MASK"):  # new Lacuna method
        lacTemp["MASK_DISTANCE_MEAN"] = tuple(numpy.abs(ptList["MASK"] * dr))
        lacTemp["MASK_DISTANCE_STD"] = tuple(numpy.abs(ptList["MASK_STD"] * dr))
    if ptList.has_key("MASK_WX"):
        lacTemp["MASK_GRAD"] = tuple(ptList["MASK"])
        lacTemp["MASK_DISTANCE_STD"] = tuple(ptList["MASK_STD"])
    if ptList.has_key("SHELL_ABSORPTION"):
        lacTemp["SHELL_ABSORPTION"] = tuple(ptList["SHELL_ABSORPTION"])
    if ptList.has_key("SHELL_ABSORPTION_STD"):
        lacTemp["SHELL_ABSORPTION_STD"] = tuple(ptList["SHELL_ABSORPTION_STD"])
    else:
        lacTemp["SHELL_ABSORPTION"] = (-1,) * dbLen
        lacTemp["SHELL_ABSORPTION_STD"] = (-1,) * dbLen
    # Lining Absorption
    if ptList.has_key("LINING_ABSORPTION"):
        lacTemp["LINING_ABSORPTION"] = tuple(ptList["LINING_ABSORPTION"])
    if ptList.has_key("LINING_ABSORPTION_STD"):
        lacTemp["LINING_ABSORPTION_STD"] = tuple(ptList["LINING_ABSORPTION_STD"])
    else:
        lacTemp["LINING_ABSORPTION"] = (-1,) * dbLen
        lacTemp["LINING_ABSORPTION_STD"] = (-1,) * dbLen
    if CanalMode == 0:
        # This doesnt work since I dont save PCA1,2,3 dumb
        # Variables that scale with PCA 1,2,3 voxel size * denotes PCA1, PCA2, PCA3
        for cAx in ["PCA1", "PCA2", "PCA3"]:
            cDr = numpy.sqrt(
                (ptList[cAx + "_X"] * dx) ** 2
                + (ptList[cAx + "_Y"] * dy) ** 2
                + (ptList[cAx + "_Z"] * dz) ** 2
            )
            rcVar = "DENSITY_PROJ_" + cAx
            if ptList.has_key("DENSITY_VOLUME_PROJ_" + cAx):
                lacTemp[rcVar] = tuple(ptList[rcVar] * cDr)
            else:
                lacTemp[rcVar] = (-1,) * dbLen
        if ptList.has_key("Canal_NUMBER"):
            lacTemp["Canal_NUMBER"] = tuple(ptList["Canal_NUMBER"])
            # lacTemp['Canal_NAME']=tuple([projectTitle+'_'+filename+'_CAN_'+str(int(curCan)) for curCan in ptList['Canal_NUMBER']])
        if ptList.has_key("Canal_NUMBER_STD"):
            lacTemp["Canal_NUMBER_STD"] = tuple(ptList["Canal_NUMBER_STD"])
        else:
            lacTemp["Canal_NUMBER_STD"] = (-1,) * dbLen
        # Nearest Neighbors
        lacTemp["NEAREST_NEIGHBOR_DISTANCE"] = (-1,) * dbLen
        lacTemp["NEAREST_NEIGHBOR_ANGLE"] = (-1,) * dbLen

        if ptList.has_key("NEIGHBORS"):
            lacTemp["NEAREST_NEIGHBOR_NEIGHBORS"] = tuple(ptList["NEIGHBORS"])
        else:
            lacTemp["NEAREST_NEIGHBOR_NEIGHBORS"] = (-1,) * dbLen
        # Mask Params
        lacTemp["POS_RADIUS"] = (-1,) * dbLen
        lacTemp["MASK_RADIUS"] = (-1,) * dbLen
        lacTemp["MASK_RADIUS_MIN"] = (-1,) * dbLen
        lacTemp["MASK_RADIUS_MAX"] = (-1,) * dbLen
        lacTemp["MASK_RADIUS_MEAN"] = (-1,) * dbLen
        lacTemp["MASK_THETA"] = (-1,) * dbLen
    if ptList.has_key("THICKNESS"):
        lacTemp["THICKNESS"] = tuple(ptList["THICKNESS"])
    else:
        lacTemp["THICKNESS"] = (-1,) * dbLen
    if ptList.has_key("THICKNESS_STD"):
        lacTemp["THICKNESS_STD"] = tuple(ptList["THICKNESS_STD"])
    else:
        lacTemp["THICKNESS_STD"] = (-1,) * dbLen
    # Lacuna Density / Volume
    if ptList.has_key("DENSITY_VOLUME"):
        lacTemp["DENSITY_VOLUME"] = tuple(
            numpy.abs(ptList["DENSITY_VOLUME"] * dx * dy * dz)
        )
        lacTemp["DENSITY"] = tuple(
            numpy.abs(1 / (ptList["DENSITY_VOLUME"] * dx * dy * dz))
        )
    elif ptList.has_key("DENSITY_VOLUME_CNT"):
        lacTemp["DENSITY_VOLUME"] = tuple(
            numpy.abs(ptList["DENSITY_VOLUME_CNT"] * dx * dy * dz)
        )
        lacTemp["DENSITY"] = tuple(
            numpy.abs(1 / (ptList["DENSITY_VOLUME_CNT"] * dx * dy * dz))
        )
    else:
        lacTemp["DENSITY_VOLUME"] = (-1,) * dbLen
        lacTemp["DENSITY"] = (-1,) * dbLen
    if CanalMode == 0:
        # Lacuna Territory Shape
        lacTemp["DISPLACEMENT_MEAN"] = (-1,) * dbLen
        if ptList.has_key("NEIGHBOR_AREA"):
            lacTemp["DENSITY_VOLUME_SHELL"] = tuple(
                numpy.abs(ptList["NEIGHBOR_AREA"] * dx * dy)
            )
        elif ptList.has_key("MASK_VOLUME_SHELL_CNT"):
            ## Old Definition of Shell
            lacTemp["DENSITY_VOLUME_SHELL"] = tuple(
                numpy.abs(ptList["MASK_VOLUME_SHELL_CNT"] * dx * dy * dz)
            )
        else:
            lacTemp["DENSITY_VOLUME_SHELL"] = (-1,) * dbLen
        # Lacuna Territory that is mineralized
        if ptList.has_key("BONE_VOLUME_CNT"):
            lacTemp["DENSITY_VOLUME_BONE"] = tuple(
                numpy.abs(ptList["BONE_VOLUME_CNT"] * dx * dy * dz)
            )
        else:
            lacTemp["DENSITY_VOLUME_BONE"] = (-1,) * dbLen
        # Lacuna Territory that is part of the mask (for porosity calculations)
        if ptList.has_key("MASK_VOLUME_CNT"):
            lacTemp["DENSITY_VOLUME_MASK"] = tuple(
                numpy.abs(ptList["MASK_VOLUME_CNT"] * dx * dy * dz)
            )
        else:
            lacTemp["DENSITY_VOLUME_MASK"] = (-1,) * dbLen
        # PCA1 is a makeshift holding place for STD until the table is once again updated
        terrShapeMap = {
            "DENSITY_VOLUME_C": "DENSITY_",
            "DENSITY_VOLUME_S": "DENSITY_STD_",
        }
        for cKey in terrShapeMap.keys():
            missingKeys = False
            for cAx in ["X", "Y", "Z"]:
                # print cKey+cAx
                if ptList.has_key(cKey + cAx):
                    # print 'isch da'
                    lacTemp[terrShapeMap[cKey] + cAx] = tuple(
                        ptList[cKey + cAx] * ptList["SCALE_" + cAx]
                    )
                else:
                    if cKey == "DENSITY_VOLUME_C":
                        missingKeys = True
                    else:
                        lacTemp[terrShapeMap[cKey] + cAx] = (-1,) * dbLen
            if not missingKeys:
                if cKey == "DENSITY_VOLUME_C":
                    dispMean = numpy.sqrt(
                        ((ptList[cKey + "X"] - ptList["POS_X"]) * dx) ** 2
                        + ((ptList[cKey + "Y"] - ptList["POS_Y"]) * dy) ** 2
                        + ((ptList[cKey + "Z"] - ptList["POS_Z"]) * dz) ** 2
                    )
                    lacTemp["DISPLACEMENT_MEAN"] = tuple(dispMean)
                    lacTemp["DISPLACEMENT_X"] = tuple(
                        (ptList[cKey + "X"] - ptList["POS_X"]) * dx
                    )
                    lacTemp["DISPLACEMENT_Y"] = tuple(
                        (ptList[cKey + "Y"] - ptList["POS_Y"]) * dy
                    )
                    lacTemp["DISPLACEMENT_Z"] = tuple(
                        (ptList[cKey + "Z"] - ptList["POS_Z"]) * dz
                    )
    # Polar Coordinates Hints
    # Only really valid for Full Femur, but Lacuna angle can be useful
    mR = numpy.sqrt(
        ((ptList["POS_X"] - numpy.mean(ptList["POS_X"])) * dx) ** 2
        + ((ptList["POS_Y"] - numpy.mean(ptList["POS_Y"])) * dy) ** 2
    )
    for cPCA in [1, 2]:
        pR = numpy.sqrt(
            ptList["PCA" + str(cPCA) + "_X"] ** 2
            + ptList["PCA" + str(cPCA) + "_Y"] ** 2
        )
        pPhi = 180 / numpy.pi * numpy.arctan2(ptList["PCA" + str(cPCA) + "_Z"], pR)
        lacTemp["PCA" + str(cPCA) + "_Phi"] = tuple(pPhi)
        lacTemp["PCA" + str(cPCA) + "_Theta"] = tuple(
            180 / numpy.pi * numpy.arccos(ptList["PCA" + str(cPCA) + "_X"] / pR)
        )  # update
    # Junk Angles
    lacTemp["POS_THETA"] = (-1,) * dbLen
    lacTemp["MASK_THETA"] = (-1,) * dbLen
    lacTemp["POS_DISTANCE"] = (-1,) * dbLen
    lacTemp["NEAREST_NEIGHBOR_AVG"] = (-1,) * dbLen
    lacTemp["NEAREST_NEIGHBOR_DISTANCE"] = (-1,) * dbLen
    lacTemp["NEAREST_NEIGHBOR_DISTANCE"] = (-1,) * dbLen
    # Normalize PCA
    pcastot = dr * numpy.sqrt(
        ptList["PCA1_S"] ** 2 + ptList["PCA2_S"] ** 2 + ptList["PCA3_S"] ** 2
    )
    # lacTemp['PCAS_TOTAL']=tuple(pcastot)
    # for tz in lacTemp.keys(): print tz+' '+str(len(lacTemp[tz]))
    outKeys = lacTemp.keys()
    outArr = [lacTemp[cKey] for cKey in outKeys]
    # for cKey in outKeys: print (cKey,len(lacTemp[cKey]))
    outMat = numpy.array(outArr).swapaxes(1, 0)
    invalidRows = numpy.sum(numpy.isnan(outMat), 1)
    outMat = outMat[numpy.nonzero(invalidRows == 0)[0], :]
    globals()["Om"] = outMat
    outMat = [tuple(obj) for obj in outMat]
    return None
    return addRowsToTable(cur, tableName, outMat, entry_order=outKeys)
    print filename + " was successfully entered %05d, invalid  %03d" % (
        lacNumOffset,
        sum(invalidRows),
    )

class CaseFreeDict:
    """Dictionary, that has case-insensitive keys.

    Keys are retained in their original form
    when queried with .keys() or .items().

    Implementation: An internal dictionary maps lowercase
    keys to (key,value) pairs. All key lookups are done
    against the lowercase keys, but all methods that expose
    keys to the user retrieve the original keys."""

    def __init__(self, dict=None):
        """Create an empty dictionary, or update from 'dict'."""
        self._dict = {}
        if dict:
            self.update(dict)

    def __getitem__(self, key):
        """Retrieve the value associated with 'key' (in any case)."""
        k = key.lower()
        return self._dict[k][1]

    def __setitem__(self, key, value):
        """Associate 'value' with 'key'. If 'key' already exists, but
        in different case, it will be replaced."""
        k = key.lower()
        self._dict[k] = (key, value)

    def has_key(self, key):
        """Case insensitive test wether 'key' exists."""
        k = key.lower()
        return self._dict.has_key(k)

    def keys(self):
        """List of keys in their original case."""
        return [v[0] for v in self._dict.values()]

    def values(self):
        """List of values."""
        return [v[1] for v in self._dict.values()]

    def items(self):
        """List of (key,value) pairs."""
        return self._dict.values()

    def get(self, key, default=None):
        """Retrieve value associated with 'key' or return default value
        if 'key' doesn't exist."""
        try:
            return self[key]
        except KeyError:
            return default

    def setdefault(self, key, default):
        """If 'key' doesn't exists, associate it with the 'default' value.
        Return value associated with 'key'."""
        if not self.has_key(key):
            self[key] = default
        return self[key]

    def update(self, dict):
        """Copy (key,value) pairs from 'dict'."""
        for k, v in dict.items():
            self[k] = v

    def __repr__(self):
        """String representation of the dictionary."""
        items = ", ".join([("%r: %r" % (k, v)) for k, v in self.items()])
        return "{%s}" % items

    def __str__(self):
        """String representation of the dictionary."""
        return repr(self)


"""
The ``bifacial`` module contains functions for modeling back surface
plane-of-array irradiance under various conditions.
"""

import pandas as pd
import numpy as np

def pvfactors_timeseries(
        solar_azimuth, solar_zenith, surface_azimuth, surface_tilt,
        axis_azimuth,
        timestamps, dni, dhi, gcr, pvrow_height, pvrow_width, albedo,
        n_pvrows=3, index_observed_pvrow=1,
        rho_front_pvrow=0.03, rho_back_pvrow=0.05,
        horizon_band_angle=15.,
        run_parallel_calculations=True, n_workers_for_parallel_calcs=2):
    """
    Calculate front and back surface plane-of-array irradiance on
    a fixed tilt or single-axis tracker PV array configuration, and using
    the open-source "pvfactors" package.  pvfactors implements the model
    described in [1]_.
    Please refer to pvfactors online documentation for more details:
    https://sunpower.github.io/pvfactors/

    Parameters
    ----------
    solar_azimuth: numeric
        Sun's azimuth angles using pvlib's azimuth convention (deg)
    solar_zenith: numeric
        Sun's zenith angles (deg)
    surface_azimuth: numeric
        Azimuth angle of the front surface of the PV modules, using pvlib's
        convention (deg)
    surface_tilt: numeric
        Tilt angle of the PV modules, going from 0 to 180 (deg)
    axis_azimuth: float
        Azimuth angle of the rotation axis of the PV modules, using pvlib's
        convention (deg). This is supposed to be fixed for all timestamps.
    timestamps: datetime or DatetimeIndex
        List of simulation timestamps
    dni: numeric
        Direct normal irradiance (W/m2)
    dhi: numeric
        Diffuse horizontal irradiance (W/m2)
    gcr: float
        Ground coverage ratio of the pv array
    pvrow_height: float
        Height of the pv rows, measured at their center (m)
    pvrow_width: float
        Width of the pv rows in the considered 2D plane (m)
    albedo: float
        Ground albedo
    n_pvrows: int, default 3
        Number of PV rows to consider in the PV array
    index_observed_pvrow: int, default 1
        Index of the PV row whose incident irradiance will be returned. Indices
        of PV rows go from 0 to n_pvrows-1.
    rho_front_pvrow: float, default 0.03
        Front surface reflectivity of PV rows
    rho_back_pvrow: float, default 0.05
        Back surface reflectivity of PV rows
    horizon_band_angle: float, default 15
        Elevation angle of the sky dome's diffuse horizon band (deg)
    run_parallel_calculations: bool, default True
        pvfactors is capable of using multiprocessing. Use this flag to decide
        to run calculations in parallel (recommended) or not.
    n_workers_for_parallel_calcs: int, default 2
        Number of workers to use in the case of parallel calculations. The
        '-1' value will lead to using a value equal to the number
        of CPU's on the machine running the model.

    Returns
    -------
    front_poa_irradiance: numeric
        Calculated incident irradiance on the front surface of the PV modules
        (W/m2)
    back_poa_irradiance: numeric
        Calculated incident irradiance on the back surface of the PV modules
        (W/m2)
    df_registries: pandas DataFrame
        DataFrame containing detailed outputs of the simulation; for
        instance the shapely geometries, the irradiance components incident on
        all surfaces of the PV array (for all timestamps), etc.
        In the pvfactors documentation, this is refered to as the "surface
        registry".

    References
    ----------
    .. [1] Anoma, Marc Abou, et al. "View Factor Model and Validation for
        Bifacial PV and Diffuse Shade on Single-Axis Trackers." 44th IEEE
        Photovoltaic Specialist Conference. 2017.
    """

    # Convert pandas Series inputs (and some lists) to numpy arrays
    if isinstance(solar_azimuth, pd.Series):
        solar_azimuth = solar_azimuth.values
    elif isinstance(solar_azimuth, list):
        solar_azimuth = np.array(solar_azimuth)
    if isinstance(solar_zenith, pd.Series):
        solar_zenith = solar_zenith.values
    if isinstance(surface_azimuth, pd.Series):
        surface_azimuth = surface_azimuth.values
    elif isinstance(surface_azimuth, list):
        surface_azimuth = np.array(surface_azimuth)
    if isinstance(surface_tilt, pd.Series):
        surface_tilt = surface_tilt.values
    if isinstance(dni, pd.Series):
        dni = dni.values
    if isinstance(dhi, pd.Series):
        dhi = dhi.values
    if isinstance(solar_azimuth, list):
        solar_azimuth = np.array(solar_azimuth)

    # Import pvfactors functions for timeseries calculations.
    from pvfactors.run import (run_timeseries_engine,
                               run_parallel_engine)

    # Build up pv array configuration parameters
    pvarray_parameters = {
        'n_pvrows': n_pvrows,
        'axis_azimuth': axis_azimuth,
        'pvrow_height': pvrow_height,
        'pvrow_width': pvrow_width,
        'gcr': gcr,
        'rho_front_pvrow': rho_front_pvrow,
        'rho_back_pvrow': rho_back_pvrow,
        'horizon_band_angle': horizon_band_angle
    }

    # Run pvfactors calculations: either in parallel or serially
    if run_parallel_calculations:
        report = run_parallel_engine(
            PVFactorsReportBuilder, pvarray_parameters,
            timestamps, dni, dhi,
            solar_zenith, solar_azimuth,
            surface_tilt, surface_azimuth,
            albedo, n_processes=n_workers_for_parallel_calcs)
    else:
        report = run_timeseries_engine(
            PVFactorsReportBuilder.build, pvarray_parameters,
            timestamps, dni, dhi,
            solar_zenith, solar_azimuth,
            surface_tilt, surface_azimuth,
            albedo)

    # Turn report into dataframe
    df_report = pd.DataFrame(report, index=timestamps)

    return df_report.total_inc_front, df_report.total_inc_back

class PVFactorsReportBuilder(object):
    """In pvfactors, a class is required to build reports when running
    calculations with multiprocessing because of python constraints"""

    @staticmethod
    def build(report, pvarray):
        """Reports will have total incident irradiance on front and
        back surface of center pvrow (index=1)"""
        # Initialize the report as a dictionary
        if report is None:
            report = {'total_inc_back': [], 'total_inc_front': []}
        # Add elements to the report
        if pvarray is not None:
            pvrow = pvarray.pvrows[1]  # use center pvrow
            report['total_inc_back'].append(
                pvrow.back.get_param_weighted('qinc'))
            report['total_inc_front'].append(
                pvrow.front.get_param_weighted('qinc'))
        else:
            # No calculation is performed when the sun is down
            report['total_inc_back'].append(np.nan)
            report['total_inc_front'].append(np.nan)

        return report

    @staticmethod
    def merge(reports):
        """Works for dictionary reports. Merges the reports list of
        dictionaries in a single dictionary. The list of the first
        dictionary are extended by those of all subsequent lists."""
        report = reports[0]
        keys_report = list(report.keys())
        for other_report in reports[1:]:  # loop won't run if len(reports) < 2
            for key in keys_report:
                report[key] += other_report[key]
        return report


"""
    Code for collecting failure trajectories using Bayesian Optimization
    Project : Policy correction using Bayesian Optimization
    Description : The file contains functions for computing failure trajectories given RL policy and
    safety specifications
"""

import numpy as np
import gym
import GPyOpt
from numpy.random import seed
from eval_policy import choose_best_action
import gym
from network import FeedForwardActorNN
import torch
import pickle
from numpy import arange
from numpy.random import rand

'''
    Bayesian Optimization module for uncovering failure trajectories

    Safety Requirement
    # Requirement 1: We would like the cartpole to not travel more than a certain
    # distance from its original location(2.4)
    # Always stay within the region (-2.4, 2.4)
'''

#=============================================Global Variables =================================#
policy = None
env = None
traj_spec_dic = {}
traj_count = 0

'''
    The function called from within the bayesian optimization module
    parameters : bounds containing the sampled variables of the state vector
    return : calls specification function and computes and returns the minimum value
'''
def sample_trajectory(sample_1,sample_2,sample_3,sample_4):
    global policy, env, traj_spec_dic,traj_count
    x1 = sample_1
    x2 = sample_2
    x3 = sample_3
    x4 = sample_4
    obs = np.array([x1,x2,x3,x4])
    #print(f'obs =========== {obs}')
    max_steps = 400
    env.reset()
    env.env.state = obs
    traj = [obs]
    actions = []
    reward = 0
    iters= 0
    ep_ret = 0
    done = False
    for _ in range(max_steps):
        iters+=1
        action = choose_best_action(obs,policy)
        actions.append(action)
        obs, rew, done, _ = env.step(action)
        #add the observation state to the current trajectory
        traj.append(obs)
        ep_ret += rew
        if done:
            break
    additional_data = {'reward':ep_ret, 'mass':env.env.total_mass}
    #Create trajectory to be sent to safety specification
    traj = (traj, additional_data)
    #print(f'trajectory ========== {traj}')
    specification_evaluation = safety_spec(traj)
    #Store the set of trajectories with negative evaluation
    if specification_evaluation<0:
        traj_spec_dic[traj_count] = (traj[0],specification_evaluation)
        traj_count = traj_count + 1
    print(f'specification_evaluation ========== {specification_evaluation}')
    return specification_evaluation

def run_Random():
    x1_max = 0.05
    x1_min = -0.05
    x2_max = 0.05
    x2_min = -0.05
    x3_max = 0.05
    x3_min = -0.05
    x4_max = 0.05
    x4_min = -0.05
    # generate a random sample from the domain
    sample_1 = x1_min + rand(1000) * (x1_max - x1_min)
    sample_2 = x2_min + rand(1000) * (x2_max - x2_min)
    sample_3 = x3_min + rand(1000) * (x3_max - x3_min)
    sample_4 = x4_min + rand(1000) * (x4_max - x4_min)
    sample = list()
    step = 0.1
    print(f'sample length ========== {len(sample_1)}')
    for i in range(len(sample_1)):
        val = sample_trajectory(sample_1[i],sample_2[i],sample_3[i],sample_4[i])
        print(f'sample1 =========== {sample_1[i]} ======== sample2 ==== {sample_2[i]} ==== sample3 ===== {sample_3[i]}')
    '''sample = list()
    step = 0.04
    for sample_1 in arange(x1_min, x1_max+step, step):
        for sample_2 in arange(x2_min, x2_max+step, step):
            for sample_3 in arange(x3_min, x3_max+step, step):
                for sample_4 in arange(x4_min, x4_max+step, step):
                    sample.append([sample_1,sample_2,sample_3,sample_4])
    print(f'sample length ========== {len(sample)}')
    for i in range(len(sample)):
        val = sample_trajectory(sample[i][0],sample[i][1],sample[i][2],sample[i][3])
        print(f'sample1 =========== {sample[i][0]} ======== sample2 ==== {sample[i][1]} ==== sample3 ===== {sample[i][2]}')'''

# 1. Always stay within the region (-2.4, 2.4)
def safety_spec(traj):
    traj = traj[0]
    #print(f'traj ========== {traj}')
    x_s = np.array(traj).T[0]
    #print(f'min value ========== {x_s}')
    return min(2.4 - np.abs(x_s))

# 2. Maintain a momentum >=-2.0 and <= 2.0
def safet_spec_2(traj):
    traj_ = traj[0]
    #print(f'traj ========== {traj}')
    mass = traj[1]['mass']
    v_s = np.array(traj_).T[1]
    return min(2. - np.abs(mass*v_s))

# 3. The angle made by the cartpole should <=0.2 within the rest position
def safet_spec_3(traj):
    traj = traj[0]
    theta=np.array(traj).T[2]
    #print(f'theta ========== {theta}')
    return min(0.2 - np.abs(theta))

if __name__ == '__main__':
    env = gym.make('CartPole-v0')
    seed = 0
    env.seed(seed)
    actor_model = 'Policies/ppo_actor_updatedCartPole-v0.pth'
    # Extract out dimensions of observation and action spaces
    obs_dim = env.observation_space.shape[0]
    act_dim = env.action_space.n #env.action_space.shape[0]

    # Build our policy the same way we build our actor model in PPO
    policy = FeedForwardActorNN(obs_dim, act_dim,True)

    # Load in the actor model saved by the PPO algorithm
    policy.load_state_dict(torch.load(actor_model))
    run_Random()
    print(f'Length trajectory ========== {len(traj_spec_dic)}')
    with open('failure_trajectory_cartpole.data', 'wb') as filehandle1:
        # store the observation data as binary data stream
        pickle.dump(traj_spec_dic, filehandle1)

Policy Refinement Using Bayesian Optimization/CartpoleRandom.py


# VIRANET TOOLS #
###############################################################################
# objSto
# Strips stoichiometric coefficients from the objectives reactions
# usage
# Inputs:
# HVM               Integrated host-virus model
# HostRxn           Host objective reaction, either:
#                   - Index value of reaction in Model.reactions        [int]
#                   - Reaction ID of the host-objective reaction        [str]
#
# Outputs:
# hostSto           Stoichiometric coefficients with metabolites for host objective
# virusSto          Stoichiometric coefficients with metabolites for virus objective
#

def objSto(HVM,HostRxn):
    "Differential usage of amino acids and nucleotides"
    # [1] Initial Setup
    # Function Dependencies
    import pandas   as pd
    import numpy    as np
    # Identify the host objective reaction
    try:
        intTest = int(HostRxn)
        hostIdx = HostRxn
    except:
        for ii in range(len(HVM.reactions)):
            if HostRxn in str(HVM.reactions[ii]):
                hostIdx = ii
    virusIdx    = len(HVM.reactions) - 1    # ViraNet(c) appends virus reaction to end in genHVM.py
    objIdx      = [hostIdx,virusIdx]
    hostID      = HVM.reactions[hostIdx].id
    virusID     = HVM.reactions[virusIdx].id
    # Condition: ensure virus reaction is virus objective
    if '_prodrxn_VN' not in virusID:
        raise ValueError('Unsupported objective, unable to analyse: refer to README')
    # Convert model into an array
    m       = HVM.to_array_based_model()
    # Create data frame
    mFrame  = pd.DataFrame(
        data    = m.S.todense(),
        columns = m.reactions.list_attr("id"),
        index   = m.metabolites.list_attr("id")
    )
    # [2] Stoichiometric Data
    # Strip the host and virus objective function stoichiometric coefficients
    hostS   = mFrame[hostID]
    virusS  = mFrame[virusID]
    # Record only non zero values
    hostSto     = hostS.loc[~(hostS==0)]
    virusSto    = virusS.loc[~(virusS==0)]
    # [3] Create output
    hostSto     = pd.DataFrame(hostSto)
    virusSto    = pd.DataFrame(virusSto)
    return (hostSto,virusSto)
###############################################################################

###############################################################################
# boundCheck
# Checks a model for any arbitrarily large bounds, removes these
# and replaces with infinite bounds
# Work based upon Kelk et al 2012
# Kelk, S. M., Olivier, B. G., Stougie, L., & Bruggeman, F. J. (2012). Optimal flux spaces of genome-scale stoichiometric models are determined by a few subnetworks. Scientific Reports, 2, 580. http://doi.org/10.1038/srep00580
# Inputs:
# Model         User-supplied model file   [.mat,.xml]
#
# Outputs:
# infModel      Model with bound corrections applied
def boundCheck(Model):
    "Reaction bound checker"
    # [1] Initial Setup
    # Function Dependencies
    import numpy    as np
    # Create pointer
    altModel = Model
    # [2] Identify and correct arbitrarily large reaction bounds
    for ii in range(len(altModel.reactions)):
        # Temporarily record the lower and upper bounds
        tmpLb   = altModel.reactions[ii].lower_bound
        tmpUb   = altModel.reactions[ii].upper_bound
        # Conditional statement
        if tmpLb <= -1000:
            altLb   = -np.inf
            altModel.reactions[ii].lower_bound = altLb
        if tmpUb >= 1000:
            altUb   = np.inf
            altModel.reactions[ii].upper_bound = altUb
    # [3] Output model
    return altModel
###############################################################################
# rangeCalculator
# Calculates , using FVA results for host and virus, the flux range to use in the
# host-derived enforcement analysis
# Inputs:
# hostIdx           Index (model.reactions) for the host-objective reaction
# virusIdx          Index (model.reactions) for the virus-objective reaction
# Optional Inputs
# solver            Declare solver to use for cobrapy: default is cglpk
# Outputs
# enfVirus          Vector of virus optima values with additional host-constraint
# maxEnfBound       Maximum enf bound
# minEnfBound       Minimum enf bound
def rangeCalculator(HVM,hostIdx,virusIdx,solver):
    "Enforcement bound creator"
    # [1] Initial Setup
    # Function Dependencies
    import cobra
    import numpy    as np
    import pandas   as pd
    # Create the virus optima vector
    enfVirus    = np.zeros((len(HVM.reactions),1))
    maxEnfBound = np.zeros((len(HVM.reactions),1))
    minEnfBound = np.zeros((len(HVM.reactions),1))
    # [2] Perform FVA for each objective
    # Objective reactions
    hostObj     = HVM.reactions[hostIdx]
    virusObj    = HVM.reactions[virusIdx]
    # Host Optimisation
    HVM.change_objective(hostObj)
    # Ensure no flux can go through the virus reaction
    # Store the bounds
    virusLb     = HVM.reactions[virusIdx].lower_bound
    virusUb     = HVM.reactions[virusIdx].upper_bound
    HVM.reactions[virusIdx].lower_bound = 0
    HVM.reactions[virusIdx].upper_bound = 0
    # FVA
    varHost     = cobra.flux_analysis.flux_variability_analysis(HVM,solver=solver)
    # Return virus objective bounds
    HVM.reactions[virusIdx].lower_bound = virusLb
    HVM.reactions[virusIdx].upper_bound = virusUb
    # Virus Optimisation
    HVM.change_objective(virusObj)
    # Ensure no flux can go through the host reaction
    # Store the bounds
    hostLb      = HVM.reactions[hostIdx].lower_bound
    hostUb      = HVM.reactions[hostIdx].upper_bound
    HVM.reactions[hostIdx].lower_bound = 0
    HVM.reactions[hostIdx].upper_bound = 0
    # FVA
    varVirus    = cobra.flux_analysis.flux_variability_analysis(HVM,solver=solver)
    # Return host objective bounds
    HVM.reactions[hostIdx].lower_bound = hostLb
    HVM.reactions[hostIdx].upper_bound = hostUb
    # Create data frames
    hostFVA     = pd.DataFrame.from_dict(varHost)
    hostFVA     = hostFVA.transpose()
    virusFVA    = pd.DataFrame.from_dict(varVirus)
    virusFVA    = virusFVA.transpose()
    # [3] Condition statements to determine the calculation and FVA steps
    # Initiate loop
    for ii in range(len(HVM.reactions)):
        # Create temporary host and virus max|min variables
        hostMax     = varHost[HVM.reactions[ii].id]['maximum']
        hostMin     = varHost[HVM.reactions[ii].id]['minimum']
        virusMax    = varVirus[HVM.reactions[ii].id]['maximum']
        virusMin    = varVirus[HVM.reactions[ii].id]['minimum']
        # Record the upper and lower bounds
        tmpLb       = HVM.reactions[ii].lower_bound
        tmpUb       = HVM.reactions[ii].upper_bound
        # Conditional: H+ > V+ && H- < V-
        if (hostMax > virusMax) and (hostMin < virusMin):
            #####################################################################
            # Calculation of bounds for condition [1]
            enfMax1 = hostMax
            enfMin1 = (hostMax - ((hostMax - virusMax) / 2))
            # Apply bounds to reaction
            HVM.reactions[ii].lower_bound   = enfMin1
            HVM.reactions[ii].upper_bound   = enfMax1
            # Zero-bound host
            HVM.reactions[hostIdx].lower_bound = 0
            HVM.reactions[hostIdx].upper_bound = 0
            # Optimize for virus
            HVM.change_objective(virusObj)
            sol = HVM.optimize(objective_sense='maximize',solver=solver)
            # Record the optima
            zMax    = sol.f
            # Return host bounds to original
            HVM.reactions[hostIdx].lower_bound = hostLb
            HVM.reactions[hostIdx].upper_bound = hostUb
            # Return reaction bounds to original
            HVM.reactions[ii].lower_bound   = tmpLb
            HVM.reactions[ii].upper_bound   = tmpUb
            # Calculation of bounds for condition [2]
            enfMin2 = hostMin
            enfMax2 = (hostMin - ((hostMin - virusMin) / 2))
            # Apply bounds to reaction
            HVM.reactions[ii].lower_bound   = enfMin2
            HVM.reactions[ii].upper_bound   = enfMax2
            # Zero-bound host
            HVM.reactions[hostIdx].lower_bound = 0
            HVM.reactions[hostIdx].upper_bound = 0
            # Optimize for virus
            HVM.change_objective(virusObj)
            sol = HVM.optimize(objective_sense='maximize',solver=solver)
            # Record the optima
            zMin    = sol.f
            # Return host bounds to original
            HVM.reactions[hostIdx].lower_bound = hostLb
            HVM.reactions[hostIdx].upper_bound = hostUb
            # Return reaction bounds to original
            HVM.reactions[ii].lower_bound   = tmpLb
            HVM.reactions[ii].upper_bound   = tmpUb
            # COMPARISON #
            # Compare zMax and zMin to find which is smallest
            if zMax < zMin:
                enfVirus[ii] = zMax
                # Record the bound
                maxEnfBound[ii] = enfMax1
                minEnfBound[ii] = enfMin1
            elif zMin < zMax:
                enfVirus[ii] = zMin
                # Record the bound
                maxEnfBound[ii] = enfMax2
                minEnfBound[ii] = enfMin2
            else:
                enfVirus[ii] = zMax
                # Record the bound
                maxEnfBound[ii] = enfMax1
                minEnfBound[ii] = enfMin1
        # Conditional: H+ > V+
        else:
            if hostMax > virusMax:
                #####################################################################
                # Calculation of bounds
                enfMax  = hostMax
                enfMin  = (hostMax - ((hostMax - virusMax) / 2))
                # Apply bounds to reaction
                HVM.reactions[ii].lower_bound   = enfMin
                HVM.reactions[ii].upper_bound   = enfMax
                # Zero-bound host
                HVM.reactions[hostIdx].lower_bound = 0
                HVM.reactions[hostIdx].upper_bound = 0
                # Optimize for virus
                HVM.change_objective(virusObj)
                sol = HVM.optimize(objective_sense='maximize',solver=solver)
                # Record the optima
                enfVirus[ii]    = sol.f
                # Return host bounds to original
                HVM.reactions[hostIdx].lower_bound = hostLb
                HVM.reactions[hostIdx].upper_bound = hostUb
                # Return reaction bounds to original
                HVM.reactions[ii].lower_bound   = tmpLb
                HVM.reactions[ii].upper_bound   = tmpUb
                # Record the bound
                maxEnfBound[ii] = enfMax
                minEnfBound[ii] = enfMin
                #####################################################################
            else:
                # Conditional: H- < V-
                if hostMin < virusMin:
                    #####################################################################
                    enfMin  = hostMin
                    enfMax  = (hostMin - ((hostMin - virusMin) / 2))
                    # Apply bounds to reaction
                    HVM.reactions[ii].lower_bound   = enfMin
                    HVM.reactions[ii].upper_bound   = enfMax
                    # Zero-bound host
                    HVM.reactions[hostIdx].lower_bound = 0
                    HVM.reactions[hostIdx].upper_bound = 0
                    # Optimize for virus
                    HVM.change_objective(virusObj)
                    sol = HVM.optimize(objective_sense='maximize',solver=solver)
                    # Record the optima
                    enfVirus[ii]    = sol.f
                    # Return host bounds to original
                    HVM.reactions[hostIdx].lower_bound = hostLb
                    HVM.reactions[hostIdx].upper_bound = hostUb
                    # Return reaction bounds to original
                    HVM.reactions[ii].lower_bound   = tmpLb
                    HVM.reactions[ii].upper_bound   = tmpUb
                    # Record the bound
                    maxEnfBound[ii] = enfMax
                    minEnfBound[ii] = enfMin
                    #####################################################################
                else:
                    #####################################################################
                    enfMax  = hostMax
                    enfMin  = hostMin
                    # Apply bounds to reaction
                    HVM.reactions[ii].lower_bound   = enfMin
                    HVM.reactions[ii].upper_bound   = enfMax
                    # Zero-bound host
                    HVM.reactions[hostIdx].lower_bound = 0
                    HVM.reactions[hostIdx].upper_bound = 0
                    # Optimize for virus
                    HVM.change_objective(virusObj)
                    sol = HVM.optimize(objective_sense='maximize',solver=solver)
                    # Record the optima
                    enfVirus[ii]    = sol.f
                    # Return host bounds to original
                    HVM.reactions[hostIdx].lower_bound = hostLb
                    HVM.reactions[hostIdx].upper_bound = hostUb
                    # Return reaction bounds to original
                    HVM.reactions[ii].lower_bound   = tmpLb
                    HVM.reactions[ii].upper_bound   = tmpUb
                    # Record the bound
                    maxEnfBound[ii] = enfMax
                    minEnfBound[ii] = enfMin
                    #####################################################################
    # [4] Output
    return (enfVirus,maxEnfBound,minEnfBound)
###############################################################################


#!/usr/bin/env python3
"""
Generate Doppler Centroid product from L0B data
"""
import os
import time
import argparse as argp
import numpy as np
from datetime import datetime

from nisar.workflows import doppler_lut_from_raw
from nisar.workflows.doppler_lut_from_raw import set_logger
from nisar.products.readers.Raw import open_rrsd
from isce3.core import TimeDelta, Linspace
from nisar.products.readers.antenna import AntennaParser

def cmd_line_parser():
    """Parse command line input arguments.

    Notes
    -----
    It also allows parsing arguments via an ASCII file
    by using prefix char "@".

    Returns
    -------
    argparse.Namespace

    """
    prs = argp.ArgumentParser(
        description='Estimate Doppler centroid from L0B raw echo and creates '
        'a 2-D Doppler LUT dumped into a CSV file',
        fromfile_prefix_chars="@",
        formatter_class=argp.ArgumentDefaultsHelpFormatter
    )
    prs.add_argument('filename_l0b', type=str,
                     help='Filename of HDF5 L0B product')
    prs.add_argument('-antenna_file', type=str, dest='antenna_file',
                     help='Filename of HDF5 Antenna product used to extract '
                     'averaged azimuth angle for EL cuts of TX + RX pol of '
                     'first beam. If not provided, the azimuth angle is '
                     'assumed to be zero!')
    prs.add_argument('-f', '--freq', type=str, choices=['A', 'B'], default='A',
                     dest='freq_band', help='Frequency band such as "A".')
    prs.add_argument('-p', '--pol', type=str, dest='txrx_pol',
                     choices=["HH", "VV", "HV", "VH"],
                     help='TxRx Polarization such as "HH". Default is the '
                     'first pol in the specified frequency band')
    prs.add_argument('-r', '--rgb', type=int, dest='num_rgb_avg', default=16,
                     help='Number of range bins to be averaged in Doppler '
                     'Estimator block. Shall be equal or larger than 1.')
    prs.add_argument('-a', '--az_block_dur', type=float, dest='az_block_dur',
                     default=4.0,
                     help='Azimuth block duration in seconds defining time-'
                     'domain correlator length used in Doppler estimator.')
    prs.add_argument('-t', '--time_interval', type=float, dest='time_interval',
                     default=2.0,
                     help='Time stamp interval between azimuth blocks in '
                     'seconds. Must not be larger than "az_block_dur".')
    prs.add_argument('-m', '--method', type=str, dest='dop_method',
                     default='CDE', choices=['SDE', 'CDE'],
                     help='Time-domain Doppler estimator methods "CDE"/"SDE"'
                     ' which are Correlator/Sign Doppler Estimator.')
    prs.add_argument('--subband', action='store_true', dest='subband',
                     help='Perform fast-time frequency subbanding on top of '
                     'time-domain correlator in Doppler estimator')
    prs.add_argument('-d', '--deg', type=int, dest='polyfit_deg',
                     default=3, help='Degree of the polyfit.')
    prs.add_argument('--polyfit', action='store_true', dest='polyfit',
                     help='If set, it will replace actual estimated doppler '
                     'by its polyfitted ones in slant range.')
    prs.add_argument('--plot', action='store_true', dest='plot',
                     help='Plot Doppler centroids and save them in '
                     '*.png files at the specified output path')
    prs.add_argument('-o', '--out', type=str, dest='out_path', default='.',
                     help='Output directory to dump Doppler product as well as'
                     'PNG plots.')

    return prs.parse_args()

def gen_doppler_range_product(args):
    """Generate Doppler-Range LUT Product.

    It generates Doppler centroid LUT as a function of slant range
    at various azimuth/pulse times and dump them into a CSV file.

    The format of the file and output filename convention is defined
    in reference [1]_.

    Parameters
    ----------
    args : argparse.Namespace
        All input arguments parsed from a command line or an ASCII file.

    References
    ----------
    .. [1] D. Kannapan, "D&C Radar Data Product SIS," JPL D-104976,
        December 3, 2020.

    """
    # Const
    PREFIX_NAME_CSV = 'NISAR_ANC'

    tic = time.time()
    # set logger
    logger = set_logger("DopplerRangeProduct")

    # get keyword args for function "doppler_lut_from_raw"
    kwargs = {key: val for key, val in args.__dict__.items() if
              'file' not in key}

    # get Raw object
    raw_obj = open_rrsd(args.filename_l0b)
    # get the SAR band char
    sar_band_char = raw_obj.sarBand
    logger.info(f'SAR band char -> {sar_band_char}')

    # operation mode, whether DBF (single or a composite channel) or
    # 'DM2' (multi-channel)
    # currently, the nunderlying module simply support DBF or single channel
    op_mode = 'DBF'

    # generate Doppler LUT2d from Raw L0B
    dop_lut, ref_utc, mask_rgb, corr_coef, txrx_pol, centerfreq, _ = \
        doppler_lut_from_raw(raw_obj, logger=logger, **kwargs)

    # check out antenna file to extract azimuth angle for EL cuts used for
    # Doppler CSV product
    if args.antenna_file is None:
        az_ang_deg = 0.0
        logger.warning(
            'No antenna file! Azimuth angle for Doppler product is '
            'assumed to be zero!'
        )
    else:
        logger.info(
            'Extracting the azimuth angle of EL cuts from antenna file.')
        ant_obj = AntennaParser(args.antenna_file)

        ant_tx = ant_obj.el_cut(pol=txrx_pol[0])
        az_ang = ant_tx.cut_angle
        # if RX pol is different from TX pol then take average of both
        if txrx_pol[0] != txrx_pol[1]:
            ant_rx = ant_obj.el_cut(pol=txrx_pol[1])
            az_ang += ant_rx.cut_angle
            az_ang *= 0.5
        az_ang_deg = np.rad2deg(az_ang)
        logger.info(
            'Azimuth angle extracted from antenna file -> '
            f'{az_ang_deg:.3f} (deg)'
        )

    # form Linspace object for uniformly-spaced azimuth time and slant range
    azt_lsp = Linspace(dop_lut.y_start, dop_lut.y_spacing, dop_lut.length)
    sr_lsp = Linspace(dop_lut.x_start, dop_lut.x_spacing, dop_lut.width)

    # get the first and last utc azimuth time w/o fractional seconds
    # in "%Y%m%dT%H%M%S" format to be used as part of CSV product filename.
    dt_utc_start = sec2str(ref_utc, azt_lsp.first)
    dt_utc_stop = sec2str(ref_utc, azt_lsp.last)
    # get current time w/o fractional seconds in "%Y%m%dT%H%M%S" format
    # used as part of CSV product filename
    dt_utc_cur = datetime.now().strftime('%Y%m%dT%H%M%S')

    # naming convention of CSV file and product spec is defined in Doc:
    # See reference [1]
    name_csv = (f'{PREFIX_NAME_CSV}_{sar_band_char}_{op_mode}_DOPP_'
                f'{dt_utc_cur}_{dt_utc_start}_{dt_utc_stop}.csv')
    file_csv = os.path.join(args.out_path, name_csv)
    logger.info(f'Dump Doppler product in "CSV" format to file -> {file_csv}')

    with open(file_csv, 'wt') as fid_csv:
        fid_csv.write(
            'UTC Time,Frequency (Hz),Doppler (Hz),Range (m),Azimuth (deg),'
            'Correlation\n'
        )
        # loop over azimuth time and slant ranges
        for i_row, azt in enumerate(azt_lsp):
            tm_utc_str = sec2isofmt(ref_utc, azt)

            for i_col, sr in enumerate(sr_lsp):
                fid_csv.write(
                    '{:s},{:.1f},{:.3f},{:.3f},{:.3f},{:.3f}\n'.format(
                        tm_utc_str, centerfreq, dop_lut.data[i_row, i_col],
                        sr, az_ang_deg,
                        mask_rgb[i_row, i_col] * corr_coef[i_row, i_col])
                    )

    # total elapsed time
    logger.info(f'Elapsed time -> {time.time() - tic:.1f} (sec)')

def sec2isofmt(ref_utc: 'isce3.core.DateTime', seconds: float) -> str:
    """seconds to isoformat string"""
    return (ref_utc + TimeDelta(seconds)).isoformat()

def sec2str(ref_utc: 'isce3.core.DateTime', seconds: float) -> str:
    """seconds to string format '%Y%m%dT%H%M%S'"""
    fmt = '%Y%m%dT%H%M%S'
    dt_iso = sec2isofmt(ref_utc, seconds)
    return datetime.fromisoformat(dt_iso.split('.')[0]).strftime(fmt)

if __name__ == "__main__":
    """Main driver"""
    gen_doppler_range_product(cmd_line_parser())


#! /usr/bin/env python
"""
MIMAS - The Multi-resolution Image Mask for Depexo Software

TODO: Write an in/out reader for MOC formats described by
http://arxiv.org/abs/1505.02937
"""

from __future__ import print_function

import logging
import numpy as np

import os
import re
from astropy.coordinates import Angle, SkyCoord
import astropy.units as u
from astropy.io import fits as pyfits
from astropy.wcs import wcs as pywcs
import healpy as hp
from .regions import Region
from .catalogs import load_table, write_table

__author__ = "Paul Hancock"
__version__ = 'v1.3.1'
__date__ = '2018-08-29'

# globals
filewcs = None

class Dummy():
    """
    A state storage class for MIMAS to work with.

    Attributes
    ----------
    add_region : list
        List of :class:`depexoTools.MIMAS.Region` to be added.

    rem_region : list
        List of :class:`depexoTools.MIMAS.Region` to be subtracted.

    include_circles : [[ra, dec, radius],...]
        List of circles to be added to the region, units are degrees.

    exclude_circles : [[ra, dec, radius], ...]
        List of circles to be subtracted from the region, units are degrees.

    include_polygons : [[ra,dec, ...], ...]
        List of polygons to be added to the region, units are degrees.

    exclude_polygons : [[ra,dec, ...], ...]
        List of polygons to be subtracted from the region, units are degrees.

    maxdepth : int
        Depth or resolution of the region for HEALPix.
        There are 4*2**maxdepth pixels at the deepest layer.
        Default = 8.

    galactic: bool
        If true then all ra/dec coordinates will be interpreted as if they were in galactic
        lat/lon (degrees)
    """
    def __init__(self, maxdepth=8):
        self.add_region = []
        self.rem_region = []
        self.include_circles = []
        self.exclude_circles = []
        self.include_polygons = []
        self.exclude_polygons = []
        self.maxdepth = maxdepth
        self.galactic = False
        return

def galactic2fk5(l, b):
    """
    Convert galactic l/b to fk5 ra/dec

    Parameters
    ----------
    l, b : float
        Galactic coordinates in radians.

    Returns
    -------
    ra, dec : float
        FK5 ecliptic coordinates in radians.
    """
    a = SkyCoord(l, b, unit=(u.radian, u.radian), frame='galactic')
    return a.fk5.ra.radian, a.fk5.dec.radian

def mask_plane(data, wcs, region, negate=False):
    """
    Mask a 2d image (data) such that pixels within 'region' are set to nan.

    Parameters
    ----------
    data : 2d-array
        Image array.

    wcs : astropy.wcs.WCS
        WCS for the image in question.

    region : :class:`depexoTools.regions.Region`
        A region within which the image pixels will be masked.

    negate : bool
        If True then pixels *outside* the region are masked.
        Default = False.

    Returns
    -------
    masked : 2d-array
        The original array, but masked as required.
    """
    # create an array but don't set the values (they are random)
    indexes = np.empty((data.shape[0]*data.shape[1], 2), dtype=int)
    # since I know exactly what the index array needs to look like i can construct
    # it faster than list comprehension would allow
    # we do this only once and then recycle it
    idx = np.array([(j, 0) for j in range(data.shape[1])])
    j = data.shape[1]
    for i in range(data.shape[0]):
        idx[:, 1] = i
        indexes[i*j:(i+1)*j] = idx

    # put ALL the pixles into our vectorized functions and minimise our overheads
    ra, dec = wcs.wcs_pix2world(indexes, 1).transpose()
    bigmask = region.sky_within(ra, dec, degin=True)
    if not negate:
        bigmask = np.bitwise_not(bigmask)
    # rework our 1d list into a 2d array
    bigmask = bigmask.reshape(data.shape)
    # and apply the mask
    data[bigmask] = np.nan
    return data

def mask_file(regionfile, infile, outfile, negate=False):
    """
    Created a masked version of file, using a region.

    Parameters
    ----------
    regionfile : str
        A file which can be loaded as a :class:`depexoTools.regions.Region`.
        The image will be masked according to this region.

    infile : str
        Input FITS image.

    outfile : str
        Output FITS image.

    negate :  bool
        If True then pixels *outside* the region are masked.
        Default = False.

    See Also
    --------
    :func:`depexoTools.MIMAS.mask_plane`
    """
    # Check that the input file is accessible and then open it
    if not os.path.exists(infile): raise AssertionError("Cannot locate fits file {0}".format(infile))
    im = pyfits.open(infile)
    if not os.path.exists(regionfile): raise AssertionError("Cannot locate region file {0}".format(regionfile))
    region = Region.load(regionfile)
    try:
        wcs = pywcs.WCS(im[0].header, naxis=2)
    except:  # TODO: figure out what error is being thrown
        wcs = pywcs.WCS(str(im[0].header), naxis=2)

    if len(im[0].data.shape) > 2:
        data = np.squeeze(im[0].data)
    else:
        data = im[0].data

    print(data.shape)
    if len(data.shape) == 3:
        for plane in range(data.shape[0]):
            mask_plane(data[plane], wcs, region, negate)
    else:
        mask_plane(data, wcs, region, negate)
    im[0].data = data
    im.writeto(outfile, overwrite=True)
    logging.info("Wrote {0}".format(outfile))
    return

def mask_table(region, table, negate=False, racol='ra', deccol='dec'):
    """
    Apply a given mask (region) to the table, removing all the rows with ra/dec inside the region
    If negate=False then remove the rows with ra/dec outside the region.

    Parameters
    ----------
    region : :class:`depexoTools.regions.Region`
        Region to mask.

    table : Astropy.table.Table
        Table to be masked.

    negate :  bool
        If True then pixels *outside* the region are masked.
        Default = False.

    racol, deccol : str
        The name of the columns in `table` that should be interpreted as ra and dec.
        Default = 'ra', 'dec'

    Returns
    -------
    masked : Astropy.table.Table
        A view of the given table which has been masked.
    """
    inside = region.sky_within(table[racol], table[deccol], degin=True)
    if not negate:
        mask = np.bitwise_not(inside)
    else:
        mask = inside
    return table[mask]

def mask_catalog(regionfile, infile, outfile, negate=False, racol='ra', deccol='dec'):
    """
    Apply a region file as a mask to a catalog, removing all the rows with ra/dec inside the region
    If negate=False then remove the rows with ra/dec outside the region.

    Parameters
    ----------
    regionfile : str
        A file which can be loaded as a :class:`depexoTools.regions.Region`.
        The catalogue will be masked according to this region.

    infile : str
        Input catalogue.

    outfile : str
        Output catalogue.

    negate :  bool
        If True then pixels *outside* the region are masked.
        Default = False.

    racol, deccol : str
        The name of the columns in `table` that should be interpreted as ra and dec.
        Default = 'ra', 'dec'

    See Also
    --------
    :func:`depexoTools.MIMAS.mask_table`

    :func:`depexoTools.catalogs.load_table`
    """
    logging.info("Loading region from {0}".format(regionfile))
    region = Region.load(regionfile)
    logging.info("Loading catalog from {0}".format(infile))
    table = load_table(infile)
    masked_table = mask_table(region, table, negate=negate, racol=racol, deccol=deccol)
    write_table(masked_table, outfile)
    return

def mim2reg(mimfile, regfile):
    """
    Convert a MIMAS region (.mim) file into a DS9 region (.reg) file.

    Parameters
    ----------
    mimfile : str
        Input file in MIMAS format.

    regfile : str
        Output file.

    """
    region = Region.load(mimfile)
    region.write_reg(regfile)
    logging.info("Converted {0} -> {1}".format(mimfile, regfile))
    return

def mim2fits(mimfile, fitsfile):
    """
    Convert a MIMAS region (.mim) file into a MOC region (.fits) file.

    Parameters
    ----------
    mimfile : str
        Input file in MIMAS format.

    fitsfile : str
        Output file.
    """
    region = Region.load(mimfile)
    region.write_fits(fitsfile, moctool='MIMAS {0}-{1}'.format(__version__, __date__))
    logging.info("Converted {0} -> {1}".format(mimfile, fitsfile))
    return

def mask2mim(maskfile, mimfile, threshold=1.0, maxdepth=8):
    """
    Use a fits file as a mask to create a region file.

    Pixels in mask file that are equal or above the threshold will be included in the reigon,
    while those that are below the threshold will not.

    Parameters
    ----------
    maskfile : str
        Input file in fits format.

    mimfile : str
        Output filename

    threshold : float
        threshold value for separating include/exclude values

    maxdepth : int
        Maximum depth (resolution) of the healpix pixels

    """
    hdu = pyfits.open(maskfile)
    wcs = pywcs.WCS(hdu[0].header)

    x, y = np.where(hdu[0].data >= threshold)
    ra, dec = wcs.all_pix2world(y, x, 0)
    sky = np.radians(Region.radec2sky(ra, dec))
    vec = Region.sky2vec(sky)
    x, y, z = np.transpose(vec)
    pix = hp.vec2pix(2**maxdepth, x, y, z, nest=True)

    region = Region(maxdepth=maxdepth)
    region.add_pixels(pix, depth=maxdepth)
    region._renorm()
    save_region(region, mimfile)
    logging.info("Converted {0} -> {1}".format(maskfile, mimfile))
    return

def box2poly(line):
    """
    Convert a string that describes a box in ds9 format, into a polygon that is given by the corners of the box

    Parameters
    ----------
    line : str
        A string containing a DS9 region command for a box.

    Returns
    -------
    poly : [ra, dec, ...]
        The corners of the box in clockwise order from top left.
    """
    words = re.split('[(\s,)]', line)
    ra = words[1]
    dec = words[2]
    width = words[3]
    height = words[4]
    if ":" in ra:
        ra = Angle(ra, unit=u.hour)
    else:
        ra = Angle(ra, unit=u.degree)
    dec = Angle(dec, unit=u.degree)
    width = Angle(float(width[:-1])/2, unit=u.arcsecond)  # strip the "
    height = Angle(float(height[:-1])/2, unit=u.arcsecond)  # strip the "
    center = SkyCoord(ra, dec)
    tl = center.ra.degree+width.degree, center.dec.degree+height.degree
    tr = center.ra.degree-width.degree, center.dec.degree+height.degree
    bl = center.ra.degree+width.degree, center.dec.degree-height.degree
    br = center.ra.degree-width.degree, center.dec.degree-height.degree
    return np.ravel([tl, tr, br, bl]).tolist()

def circle2circle(line):
    """
    Parse a string that describes a circle in ds9 format.

    Parameters
    ----------
    line : str
        A string containing a DS9 region command for a circle.

    Returns
    -------
    circle : [ra, dec, radius]
        The center and radius of the circle.
    """
    words = re.split('[(,\s)]', line)
    ra = words[1]
    dec = words[2]
    radius = words[3][:-1]  # strip the "
    if ":" in ra:
        ra = Angle(ra, unit=u.hour)
    else:
        ra = Angle(ra, unit=u.degree)
    dec = Angle(dec, unit=u.degree)
    radius = Angle(radius, unit=u.arcsecond)
    return [ra.degree, dec.degree, radius.degree]

def poly2poly(line):
    """
    Parse a string of text containing a DS9 description of a polygon.

    This function works but is not very robust due to the constraints of healpy.

    Parameters
    ----------
    line : str
        A string containing a DS9 region command for a polygon.

    Returns
    -------
    poly : [ra, dec, ...]
        The coordinates of the polygon.
    """
    words = re.split('[(\s,)]', line)
    ras = np.array(words[1::2])
    decs = np.array(words[2::2])
    coords = []
    for ra, dec in zip(ras, decs):
        if ra.strip() == '' or dec.strip() == '':
            continue
        if ":" in ra:
            pos = SkyCoord(Angle(ra, unit=u.hour), Angle(dec, unit=u.degree))
        else:
            pos = SkyCoord(Angle(ra, unit=u.degree), Angle(dec, unit=u.degree))
        # only add this point if it is some distance from the previous one
        coords.extend([pos.ra.degree, pos.dec.degree])
    return coords

def reg2mim(regfile, mimfile, maxdepth):
    """
    Parse a DS9 region file and write a MIMAS region (.mim) file.

    Parameters
    ----------
    regfile : str
        DS9 region (.reg) file.

    mimfile : str
        MIMAS region (.mim) file.

    maxdepth : str
        Depth/resolution of the region file.

    """
    logging.info("Reading regions from {0}".format(regfile))
    lines = (l for l in open(regfile, 'r') if not l.startswith('#'))
    poly = []
    circles = []
    for line in lines:
        if line.startswith('box'):
            poly.append(box2poly(line))
        elif line.startswith('circle'):
            circles.append(circle2circle(line))
        elif line.startswith('polygon'):
            logging.warning("Polygons break a lot, but I'll try this one anyway.")
            poly.append(poly2poly(line))
        else:
            logging.warning("Not sure what to do with {0}".format(line[:-1]))
    container = Dummy(maxdepth=maxdepth)
    container.include_circles = circles
    container.include_polygons = poly

    region = combine_regions(container)
    save_region(region, mimfile)
    return

def combine_regions(container):
    """
    Return a region that is the combination of those specified in the container.
    The container is typically a results instance that comes from argparse.

    Order of construction is: add regions, subtract regions, add circles, subtract circles,
    add polygons, subtract polygons.

    Parameters
    ----------
    container : :class:`depexoTools.MIMAS.Dummy`
        The regions to be combined.

    Returns
    -------
    region : :class:`depexoTools.regions.Region`
        The constructed region.
    """
    # create empty region
    region = Region(container.maxdepth)

    # add/rem all the regions from files
    for r in container.add_region:
        logging.info("adding region from {0}".format(r))
        r2 = Region.load(r[0])
        region.union(r2)

    for r in container.rem_region:
        logging.info("removing region from {0}".format(r))
        r2 = Region.load(r[0])
        region.without(r2)

    # add circles
    if len(container.include_circles) > 0:
        for c in container.include_circles:
            circles = np.radians(np.array(c))
            if container.galactic:
                l, b, radii = circles.reshape(3, circles.shape[0]//3)
                ras, decs = galactic2fk5(l, b)
            else:
                ras, decs, radii = circles.reshape(3, circles.shape[0]//3)
            region.add_circles(ras, decs, radii)

    # remove circles
    if len(container.exclude_circles) > 0:
        for c in container.exclude_circles:
            r2 = Region(container.maxdepth)
            circles = np.radians(np.array(c))
            if container.galactic:
                l, b, radii = circles.reshape(3, circles.shape[0]//3)
                ras, decs = galactic2fk5(l, b)
            else:
                ras, decs, radii = circles.reshape(3, circles.shape[0]//3)
            r2.add_circles(ras, decs, radii)
            region.without(r2)

    # add polygons
    if len(container.include_polygons) > 0:
        for p in container.include_polygons:
            poly = np.radians(np.array(p))
            poly = poly.reshape((poly.shape[0]//2, 2))
            region.add_poly(poly)

    # remove polygons
    if len(container.exclude_polygons) > 0:
        for p in container.include_polygons:
            poly = np.array(np.radians(p))
            r2 = Region(container.maxdepth)
            r2.add_poly(poly)
            region.without(r2)

    return region

def intersect_regions(flist):
    """
    Construct a region which is the intersection of all regions described in the given
    list of file names.

    Parameters
    ----------
    flist : list
        A list of region filenames.

    Returns
    -------
    region : :class:`depexoTools.regions.Region`
        The intersection of all regions, possibly empty.
    """
    if len(flist) < 2:
        raise Exception("Require at least two regions to perform intersection")
    a = Region.load(flist[0])
    for b in [Region.load(f) for f in flist[1:]]:
        a.intersect(b)
    return a

def save_region(region, filename):
    """
    Save the given region to a file

    Parameters
    ----------
    region : :class:`depexoTools.regions.Region`
        A region.

    filename : str
        Output file name.
    """
    region.save(filename)
    logging.info("Wrote {0}".format(filename))
    return

def save_as_image(region, filename):
    """
    Convert a MIMAS region (.mim) file into a image (eg .png)

    Parameters
    ----------
    region : :class:`depexoTools.regions.Region`
        Region of interest.

    filename : str
        Output filename.
    """
    import healpy as hp
    pixels = list(region.get_demoted())
    order = region.maxdepth
    m = np.arange(hp.nside2npix(2**order))
    m[:] = 0
    m[pixels] = 1
    hp.write_map(filename, m, nest=True, coord='C')
    return

import numpy,import astropy,from astropy