#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Created on Tue May 23 08:46:02 2017
@author: mschull
"""

# This file is part of PyDisALEXI for running disALEXI using different TSEB models
# Copyright 2016 Mitchell Schull and contributors listed in the README.md file.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
import os
import numpy as np
import subprocess
from osgeo import gdal
from osgeo.gdalconst import GA_ReadOnly
import pandas as pd
from .TSEB_usda import TSEB_PT_usda
from .utils import writeArray2Tiff, getParFromExcel, warp, folders
from scipy import ndimage, interp
from .landsatTools import landsat_metadata, GeoTIFF
from .TSEB_utils_usda import sunset_sunrise, interp_ta
# from joblib import Parallel, delayed
from astropy.convolution import Gaussian2DKernel, Box2DKernel
from astropy.convolution import convolve_fft
from joblib import Memory

cachedir = os.path.join(os.getcwd(), 'cachedir')
if not os.path.exists(cachedir):
    os.mkdir(cachedir)

memory = Memory(cachedir, verbose=0)

def _DisALEXI_PT(ET_ALEXI,
                 Rs_1,
                 Rs24in,
                 Tr_K,
                 Ta,
                 vza,
                 u,
                 p,
                 zs,
                 aleafv,
                 aleafn,
                 aleafl,
                 adeadv,
                 adeadn,
                 adeadl,
                 albedo,
                 ndvi,
                 LAI,
                 clump,
                 hc,
                 mask,
                 time,
                 t_rise,
                 t_end,
                 leaf_width=1.,
                 alpha_PT=1.32):
    '''DisALEXI based on Priestley-Taylor TSEB

        Calculates the Priestley Taylor TSEB fluxes using a single observation of
        composite radiometric temperature and using resistances in series.

        Parameters
        ----------
        ET_ALEXI : float
            Coarse resolution daily ET from ALEXI
        geoDict : dictionary
            Dictionary containing:
            inProj4 : proj4 string
                ALEXI ET proj4 string
            outProj4 : proj4 string
                DisALEXI ET proj4 string
            inUL : float array
                Upper left lat/lon coordinates of ALEXI image
            inRes : float array
                ALEXI ET lat/lon resolution
        Rs_1 : float
            Overpass insolation (w m-2)
        Rs24 : float
            Total daily insolation (w m-2)
        Tr_K : float
            Radiometric composite temperature (Kelvin).
        vza : float
            View Zenith Angle (degrees).
        u : float
            Wind speed above the canopy (m s-1).
        ea : float
            Water vapour pressure above the canopy (mb).
        p : float
            Atmospheric pressure (mb), use 1013 mb by default.
        Sn_C : float
            Canopy net shortwave radiation (W m-2).
        Sn_S : float
            Soil net shortwave radiation (W m-2).
        LAI : float
            Effective Leaf Area Index (m2 m-2).
        hc : float
            Canopy height (m).
        emis_C : float
            Leaf emissivity.
        emis_S : flaot
            Soil emissivity.
        z_0M : float
            Aerodynamic surface roughness length for momentum transfer (m).
        d_0 : float
            Zero-plane displacement height (m).
        z_u : float
            Height of measurement of windspeed (m).
        z_T : float
            Height of measurement of air temperature (m).
        leaf_width : float, optional
            average/effective leaf width (m).
        z0_soil : float, optional
            bare soil aerodynamic roughness length (m).
        alpha_PT : float, optional
            Priestley Taylor coeffient for canopy potential transpiration,
            use 1.26 by default.
        x_LAD : float, optional
            Campbell 1990 leaf inclination distribution function chi parameter.
        f_c : float, optional
            Fractional cover.
        f_g : float, optional
            Fraction of vegetation that is green.
        w_C : float, optional
            Canopy width to height ratio.
        resistance_form : int, optional
            Flag to determine which Resistances R_x, R_S model to use.

                * 0 [Default] Norman et al 1995 and Kustas et al 1999.
                * 1 : Choudhury and Monteith 1988.
                * 2 : McNaughton and Van der Hurk 1995.

        calcG_params : list[list,float or array], optional
            Method to calculate soil heat flux,parameters.

                * [[1],G_ratio]: default, estimate G as a ratio of Rn_S, default Gratio=0.35.
                * [[0],G_constant] : Use a constant G, usually use 0 to ignore the computation of G.
                * [[2,Amplitude,phase_shift,shape],time] : estimate G from Santanello and Friedl with G_param list of parameters (see :func:`~TSEB.calc_G_time_diff`).
        UseL : float or None, optional
            If included, its value will be used to force the Moning-Obukhov stability length.

        Returns
        -------
        flag : int
            Quality flag, see Appendix for description.
        T_S : float
            Soil temperature  (Kelvin).
        T_C : float
            Canopy temperature  (Kelvin).
        T_AC : float
            Air temperature at the canopy interface (Kelvin).
        L_nS : float
            Soil net longwave radiation (W m-2)
        L_nC : float
            Canopy net longwave radiation (W m-2)
        LE_C : float
            Canopy latent heat flux (W m-2).
        H_C : float
            Canopy sensible heat flux (W m-2).
        LE_S : float
            Soil latent heat flux (W m-2).
        H_S : float
            Soil sensible heat flux (W m-2).
        G : float
            Soil heat flux (W m-2).
        R_S : float
            Soil aerodynamic resistance to heat transport (s m-1).
        R_x : float
            Bulk canopy aerodynamic resistance to heat transport (s m-1).
        R_A : float
            Aerodynamic resistance to heat transport (s m-1).
        u_friction : float
            Friction velocity (m s-1).
        L : float
            Monin-Obuhkov length (m).
        n_iterations : int
            number of iterations until convergence of L.

        References
        ----------
        .. [Norman1995] J.M. Norman, W.P. Kustas, K.S. Humes, Source approach for estimating
            soil and vegetation energy fluxes in observations of directional radiometric
            surface temperature, Agricultural and Forest Meteorology, Volume 77, Issues 3-4,
            Pages 263-293,
            http://dx.doi.org/10.1016/0168-1923(95)02265-Y.
        .. [Kustas1999] William P Kustas, John M Norman, Evaluation of soil and vegetation heat
            flux predictions using a simple two-source model with radiometric temperatures for
            partial canopy cover, Agricultural and Forest Meteorology, Volume 94, Issue 1,
            Pages 13-29,
            http://dx.doi.org/10.1016/S0168-1923(99)00005-2.
        '''

    # Set up input parameters
    MatXsize = 7
    Tr_Kresize = np.tile(np.array(np.resize(Tr_K, [np.size(Tr_K), 1])), (1, MatXsize))
    vzaresize = np.tile(np.resize(vza, [np.size(vza), 1]), (1, MatXsize))
    #        T_A_Kresize = np.tile(range(270,340,10),(np.size(vza),1))
    Tr_ADD = np.tile(np.transpose(range(0, 20, 3)), [np.size(hc), 1])
    #        Tr_Kcol = np.resize(Tr_K,[np.size(Tr_K),1])
    Tr_Kcol = np.resize(Ta, [np.size(Ta), 1])
    T_A_Kresize = Tr_Kcol + Tr_ADD
    uresize = np.tile(np.resize(u, [np.size(u), 1]), (1, MatXsize))
    presize = np.tile(np.resize(p, [np.size(p), 1]), (1, MatXsize))
    Rs_1resize = np.tile(np.resize(Rs_1, [np.size(Rs_1), 1]), (1, MatXsize))
    zsresize = np.tile(np.resize(zs, [np.size(zs), 1]), (1, MatXsize))
    aleafvresize = np.tile(np.resize(aleafv, [np.size(hc), 1]), (1, MatXsize))
    aleafnresize = np.tile(np.resize(aleafn, [np.size(hc), 1]), (1, MatXsize))
    aleaflresize = np.tile(np.resize(aleafl, [np.size(hc), 1]), (1, MatXsize))
    adeadvresize = np.tile(np.resize(adeadv, [np.size(hc), 1]), (1, MatXsize))
    adeadnresize = np.tile(np.resize(adeadn, [np.size(hc), 1]), (1, MatXsize))
    adeadlresize = np.tile(np.resize(adeadl, [np.size(hc), 1]), (1, MatXsize))
    albedoresize = np.tile(np.resize(albedo, [np.size(hc), 1]), (1, MatXsize))
    ndviresize = np.tile(np.resize(ndvi, [np.size(hc), 1]), (1, MatXsize))
    LAIresize = np.tile(np.resize(LAI, [np.size(LAI), 1]), (1, MatXsize))
    clumpresize = np.tile(np.resize(clump, [np.size(hc), 1]), (1, MatXsize))
    hcresize = np.tile(np.resize(hc, [np.size(hc), 1]), (1, MatXsize))
    maskresize = np.tile(np.array(np.resize(mask, [np.size(hc), 1])), (1, MatXsize))
    timeresize = np.tile(np.array(np.resize(time, [np.size(hc), 1])), (1, MatXsize))
    t_riseresize = np.tile(np.array(np.resize(t_rise, [np.size(hc), 1])), (1, MatXsize))
    t_endresize = np.tile(np.array(np.resize(t_end, [np.size(hc), 1])), (1, MatXsize))
    leaf_widthresize = np.tile(np.resize(leaf_width, [np.size(hc), 1]), (1, MatXsize))
    alpha_PTresize = np.tile(np.resize(alpha_PT, [np.size(hc), 1]), (1, MatXsize))

    # run TSEB over TA options
    output = TSEB_PT_usda(
        Tr_Kresize,
        vzaresize,
        T_A_Kresize,
        uresize,
        presize,
        Rs_1resize,
        zsresize,
        aleafvresize,
        aleafnresize,
        aleaflresize,
        adeadvresize,
        adeadnresize,
        adeadlresize,
        albedoresize,
        ndviresize,
        LAIresize,
        clumpresize,
        hcresize,
        maskresize,
        timeresize,
        t_riseresize,
        t_endresize,
        leaf_width=leaf_widthresize,
        a_PT_in=alpha_PTresize)

    scaling = 1.0
    Fsun = (output[4] + output[6]) / np.resize(Rs_1, [np.size(hc), 1])
    EFeq = Fsun * (np.reshape(Rs24in, [np.size(hc), 1]))
    et = EFeq / 2.45 * scaling
    et[et < 0.01] = 0.01

    # =============find Average ETd======================================

    et_alexi = np.array(np.reshape(ET_ALEXI, [np.size(ET_ALEXI)]) * 10000, dtype='int')
    etDict = {'ID': et_alexi,
              'et1': np.reshape(et[:, 0], [np.size(ET_ALEXI)]),
              'et2': np.reshape(et[:, 1], [np.size(ET_ALEXI)]),
              'et3': np.reshape(et[:, 2], [np.size(ET_ALEXI)]),
              'et4': np.reshape(et[:, 3], [np.size(ET_ALEXI)]),
              'et5': np.reshape(et[:, 4], [np.size(ET_ALEXI)]),
              'et6': np.reshape(et[:, 5], [np.size(ET_ALEXI)]),
              'et7': np.reshape(et[:, 6], [np.size(ET_ALEXI)])}

    etDF = pd.DataFrame(etDict, columns=etDict.keys())
    etDF = pd.DataFrame(etDict)
    group = etDF.groupby(etDF['ID'])
    valMean = group.mean()
    outData = np.zeros(et.shape)
    for i in range(valMean.shape[0]):
        outData[et_alexi == valMean.index[i]] = valMean.iloc[i]
    et = np.reshape(outData, et.shape)
    # ======interpolate over mutiple Ta ===================================

    from scipy.interpolate import interp1d
    x = range(0, 20, 3)
    ET_ALEXI[mask == 0] = -9999.
    et_alexi = np.reshape(ET_ALEXI, [np.size(hc), 1])
    bias = et_alexi - et
    # check if all values inrow are nan
    nanIndex = np.sum(np.isnan(bias), axis=1)
    # set all to 1 so it doesnt throw an error below
    bias[np.where(nanIndex == MatXsize), :] = 1.
    f_bias = interp1d(x, bias, kind='linear', bounds_error=False)
    f_ta = interp1d(x, T_A_Kresize, kind='linear', bounds_error=False)

    biasInterp = f_bias(np.linspace(0, 20, 1000))
    TaInterp = f_ta(np.linspace(0, 20, 1000))
    # extract the Ta based on minimum bias at Fine resolution
    minBiasIndex = np.array(np.nanargmin(abs(biasInterp), axis=1))
    TaExtrap = TaInterp[np.array(range(np.size(hc))), minBiasIndex]
    TaExtrap[np.where(nanIndex == MatXsize)] = np.nan
    Tareshape = np.reshape(TaExtrap, np.shape(hc))

    T_A_K = Tareshape
    output = {'T_A_K': T_A_K}
    return output

class disALEXI(object):
    def __init__(self, fn, dt, isUSA):
        #        base = os.path.abspath(os.path.join(fn,os.pardir,os.pardir,os.pardir,
        #                                            os.pardir,os.pardir))
        base = os.getcwd()

        Folders = folders(base)
        self.landsatSR = Folders['landsatSR']
        #        self.landsatNDVI = Folders['landsatNDVI']
        #        self.ALEXIbase = Folders['ALEXIbase']
        #        self.metBase = Folders['metBase']
        #        self.landsatDataBase = Folders['landsatDataBase']
        self.resultsBase = Folders['resultsBase']
        self.fn = fn
        self.meta = landsat_metadata(fn)
        self.sceneID = self.meta.LANDSAT_SCENE_ID
        self.productID = fn.split(os.sep)[-1][:-8]
        #        self.productID = self.meta.LANDSAT_PRODUCT_ID
        self.scene = self.sceneID[3:9]
        self.isUSA = isUSA
        self.dt = dt
        self.satscene_path = os.sep.join(fn.split(os.sep)[:-2])

    def DisALEXI_PT(self,
                    ET_ALEXI,
                    Rs_1,
                    Rs24in,
                    Tr_K,
                    Ta,
                    vza,
                    u,
                    p,
                    zs,
                    aleafv,
                    aleafn,
                    aleafl,
                    adeadv,
                    adeadn,
                    adeadl,
                    albedo,
                    ndvi,
                    LAI,
                    clump,
                    hc,
                    mask,
                    time,
                    t_rise,
                    t_end,
                    leaf_width=1.,
                    alpha_PT=1.32):
        disalexi = _DisALEXI_PT
        return disalexi(ET_ALEXI, Rs_1, Rs24in, Tr_K, Ta, vza, u, p, zs, aleafv, aleafn,
                        aleafl, adeadv, adeadn, adeadl, albedo, ndvi, LAI, clump, hc,
                        mask, time, t_rise, t_end, leaf_width=1., alpha_PT=1.32)

    def smoothTaData(self, ALEXIgeodict):

        ALEXILatRes = ALEXIgeodict['ALEXI_LatRes']
        ALEXILonRes = ALEXIgeodict['ALEXI_LonRes']
        sceneID = self.sceneID
        scene = self.scene
        outFN = os.path.join(self.resultsBase, scene, '%s_Ta.tif' % sceneID[:-5])
        inProj4 = '+proj=longlat +ellps=WGS84 +datum=WGS84 +no_defs'
        # =======================convert fine TA to coarse resolution=========
        outfile = os.path.join(self.resultsBase, scene, 'testTa_DisALEXI.tif')
        #            outfile = os.path.join(self.resultsBase,scene,'%s_Ta.tif' % sceneID[:-5])

        coarseFile = os.path.join(self.resultsBase, scene, 'TaCoarse.tif')
        coarse2fineFile = os.path.join(self.resultsBase, scene, 'TaCoarse2Fine.tif')
        #            outFN = coarseFile[:-10]+'.tif'

        if not os.path.exists(outFN):
            print 'get->Ta'
            # get mask from Landsat LAI
            ls = GeoTIFF(outfile)
            #                laiFN = os.path.join(self.landsatDataBase,'LAI',scene,'lndlai.%s.hdf' % sceneID)
            #                hdf = SD(laiFN,SDC.READ)
            #                data2D = hdf.select('cfmask')
            #                cfmask = data2D[:,:].astype(np.double)
            sceneDir = os.path.join(self.satscene_path, 'CF_MASK')
            maskFN = os.path.join(sceneDir, '%s_Mask.tiff' % sceneID)
            g = gdal.Open(maskFN, GA_ReadOnly)
            cfmask = g.ReadAsArray()
            g = None
            # =============find Average Ta====================================== COMMENTED FOR TESTING
            in_ds = gdal.Open(outfile)
            coarseds = gdal.Translate(coarseFile, in_ds,
                                      options=gdal.TranslateOptions(
                                          resampleAlg='average',
                                          xRes=400,
                                          yRes=400))
            fineds = gdal.Warp(outFN, coarseds, options=gdal.WarpOptions(resampleAlg='average',
                                                                         height=ls.nrow,
                                                                         width=ls.ncol))
            coarseds = None
            # ========smooth Ta data========================================
            ta = fineds.ReadAsArray()
            fineRes = ls.Lat[1, 0] - ls.Lat[0, 0]
            coarseRes = ALEXILatRes
            course2fineRatio = coarseRes ** 2 / fineRes ** 2
            rid2 = int(np.sqrt(course2fineRatio))
            #            gauss_kernal = Gaussian2DKernel(rid2)
            box_kernal = Box2DKernel(rid2)
            ta = convolve_fft(ta, box_kernal, allow_huge=True)
            fineds.GetRasterBand(1).WriteArray(ta)
            fineds = None

            ulx = ls.ulx
            uly = ls.uly
            delx = ls.delx
            dely = -ls.dely
            fineRes = ls.Lat[1, 0] - ls.Lat[0, 0]
            coarseRes = ALEXILatRes
            inUL = [ulx, uly]
            inRes = [delx, dely]

            #            Ta = interp_ta(ta,coarseRes,fineRes)-273.16
            Ta = ta - 273.16  # FOR TESTING!!

            outFormat = gdal.GDT_Float32
            writeArray2Tiff(Ta, inRes, inUL, ls.proj4, outFN, outFormat)
            os.remove(coarseFile)

    def runDisALEXI(self, xStart, yStart, xSize, ySize, ALEXIgeodict, TSEB_only):
        # USER INPUT============================================================
        ALEXILatRes = ALEXIgeodict['ALEXI_LatRes']
        ALEXILonRes = ALEXIgeodict['ALEXI_LonRes']
        sceneID = self.sceneID
        scene = self.scene
        productID = self.productID
        #        xSize = 200
        #        ySize = 200

        # -------------pick Landcover map----------------
        if self.isUSA == 1:
            landcover = 'NLCD'
        else:
            landcover = 'GlobeLand30'

        yeardoy = sceneID[9:16]
        # -------------get Landsat information-----------
        #        ls = GeoTIFF(os.path.join(self.landsatSR,'temp',"%s_band10.tif" % productID))
        sceneDir = os.path.join(self.satscene_path, 'LST')
        ls = GeoTIFF(os.path.join(sceneDir, '%s_lstSharp.tiff' % sceneID))
        g = gdal.Open(os.path.join(sceneDir, '%s_lstSharp.tiff' % sceneID))
        solZen = self.meta.SUN_ELEVATION
        #        nsamples = int(self.meta.REFLECTIVE_SAMPLES)
        #        nlines = int(self.meta.REFLECTIVE_LINES)
        nsamples = g.RasterXSize
        nlines = g.RasterYSize
        if xStart == ((nsamples / xSize) * xSize):
            xSize = (nsamples - xStart)
        if yStart == ((nlines / ySize) * ySize):
            ySize = (nlines - yStart)
        g = None
        inProj4 = '+proj=longlat +ellps=WGS84 +datum=WGS84 +no_defs'
        sz = np.radians(90 - solZen)  # convert sza to radians

        # ===========================get the ETd data==============================

        #        sceneDir = os.path.join(self.ALEXIbase,'%s' % scene)
        sceneDir = os.path.join(self.satscene_path, 'ET', '400m')
        outFN = os.path.join(sceneDir, '%s_alexiET.tiff' % sceneID)
        g = gdal.Open(outFN, GA_ReadOnly)
        ET_ALEXI = g.ReadAsArray(xStart, yStart, xSize, ySize)
        g = None

        # =============get MET data================================================
        # get CFSR MET data at overpass time

        #        sceneDir = os.path.join(self.metBase,'%s' % scene)
        sceneDir = os.path.join(self.satscene_path, 'MET')
        # ------------get-> surface pressure...
        outFN = os.path.join(sceneDir, '%s_p.tiff' % sceneID)
        g = gdal.Open(outFN, GA_ReadOnly)
        #        print("xsize:%d" % xSize)
        #        print("ysize:%d" % ySize)
        p = g.ReadAsArray(xStart, yStart, xSize, ySize)
        p /= 100.  # convert to mb
        g = None

        # ------------get-> ea...
        outFN = os.path.join(sceneDir, '%s_q2.tiff' % sceneID)
        g = gdal.Open(outFN, GA_ReadOnly)
        q2 = g.ReadAsArray(xStart, yStart, xSize, ySize)
        g = None

        ea = ((q2 * (1000. / 621.9907)) * (p * 100.)) * 0.001  # kPa
        ea *= 10.  # mb
        # ====get CFSR air temperature==========================================
        outFN = os.path.join(sceneDir, '%s_Ta.tiff' % sceneID)
        g = gdal.Open(outFN, GA_ReadOnly)
        Ta = g.ReadAsArray(xStart, yStart, xSize, ySize)
        g = None

        outFN = os.path.join(self.resultsBase, scene, '%s_Ta.tif' % sceneID[:-5])
        if (TSEB_only == 1):
            #            if ((xStart==0) & (yStart==0)):
            #                #ls = GeoTIFF(os.path.join(self.landsatSR, scene,'%s_sr_band1.tif' % productID))
            #
            #                #=======================convert fine TA to coarse resolution=========
            #                outfile = os.path.join(self.resultsBase,scene,'Ta_DisALEXI.tif')
            #    #            outfile = os.path.join(self.resultsBase,scene,'%s_Ta.tif' % sceneID[:-5])
            #
            #                coarseFile = os.path.join(self.resultsBase,scene,'TaCoarse.tif')
            #                coarse2fineFile = os.path.join(self.resultsBase,scene,'TaCoarse2Fine.tif')
            #    #            outFN = coarseFile[:-10]+'.tif'
            #
            #                if not os.path.exists(outFN):
            #                    print 'get->Ta'
            #                    # get mask from Landsat LAI
            #                    ls = GeoTIFF(outfile)
            #    #                laiFN = os.path.join(self.landsatDataBase,'LAI',scene,'lndlai.%s.hdf' % sceneID)
            #    #                hdf = SD(laiFN,SDC.READ)
            #    #                data2D = hdf.select('cfmask')
            #    #                cfmask = data2D[:,:].astype(np.double)
            #
            #                    maskFN = os.path.join(self.landsatDataBase,'Mask',scene,'%s_mask.tiff' % sceneID)
            #                    g = gdal.Open(maskFN,GA_ReadOnly)
            #                    cfmask = g.ReadAsArray()
            #                    g= None
            #                    g = gdal.Open(outfile,GA_ReadOnly)
            #                    ta = g.ReadAsArray()
            #                    ta[cfmask > 0]=0
            #                    g= None
            #                    mask = os.path.join(self.resultsBase,scene,"TafineMask.tif")
            #                    masked = os.path.join(self.resultsBase,scene,"TafineMasked.tif")
            #                    ls.clone(mask,ta)
            #                    subprocess.check_output('gdal_fillnodata.py %s %s -mask %s -of GTiff' % (outfile,masked,mask),shell=True)
            #                    optionList = ['-overwrite', '-s_srs', '%s' % ls.proj4,'-t_srs',
            #                                  '%s' % inProj4,'-r', 'average','-tr',
            #                                  '%f' % ALEXILatRes, '%f' % ALEXILonRes,
            #                                  '-srcnodata','270.','-dstnodata','0.0',
            #                                  '-of','GTiff','%s' % masked, '%s' % coarseFile]
            #
            #                    warp(optionList)
            #                    #=======now convert the averaged coarse Ta to fine resolution==
            #                    nrow = ls.nrow+100.
            #                    ncol = ls.ncol+100.
            #                    optionList = ['-overwrite', '-s_srs', '%s' % inProj4, '-t_srs',
            #                                  '%s' % ls.proj4,'-r', 'near','-ts',
            #                                  '%f' % nrow, '%f' % ncol,'-of',
            #                                  'GTiff','%s' % coarseFile, '%s' % coarse2fineFile]
            #                    warp(optionList)
            #                    #========smooth Ta data========================================
            #                    ulx = ls.ulx
            #                    uly = ls.uly
            #                    lrx = ls.lrx
            #                    lry = ls.lry
            #                    delx = ls.delx
            #                    dely = -ls.dely
            #                    fineRes = ls.Lat[1,0]-ls.Lat[0,0]
            #                    coarseRes = ALEXILatRes
            #                    inUL = [ulx,uly]
            #                    inRes = [delx,dely]
            #                    g = gdal.Open(coarse2fineFile,GA_ReadOnly)
            #                    ta = g.ReadAsArray()
            #                    g= None
            #
            #                    Ta = interp_ta(ta,coarseRes,fineRes)-273.16
            #
            #                    outFormat = gdal.GDT_Float32
            #                    writeArray2Tiff(Ta,inRes,inUL,ls.proj4,outFN,outFormat)
            #    #
            #    #                optionList = ['-overwrite','-te', '%f' % ulx, '%f' % lry,
            #    #                              '%f' % lrx,'%f' % uly,'-tr',
            #    #                              '%f' % delx, '%f' % dely ,'-multi','-of','GTiff',
            #    #                              '%s' % coarse2fineFile, '%s' % outFN]
            #    #                warp(optionList)
            #                    #os.remove(coarseFile)
            g = gdal.Open(outFN, GA_ReadOnly)
            T_A_K = g.ReadAsArray(xStart, yStart, xSize, ySize) + 273.16
            g = None

        #        sceneDir = os.path.join(self.metBase,'%s' % scene)
        outFN = os.path.join(sceneDir, '%s_u.tiff' % sceneID)
        g = gdal.Open(outFN, GA_ReadOnly)
        u = g.ReadAsArray(xStart, yStart, xSize, ySize)
        g = None
        # ===create lat long files==============================================

        inUL = [ls.ulx, ls.uly]
        inRes = [ls.delx, -ls.dely]
        lat_fName = os.path.join(self.landsatSR, 'temp', 'lat.tif')
        lon_fName = os.path.join(self.landsatSR, 'temp', 'lon.tif')
        if not os.path.exists(lat_fName):
            lats = ls.Lat_pxcenter
            lons = ls.Lon_pxcenter
            writeArray2Tiff(lats, inRes, inUL, ls.proj4, lat_fName, gdal.GDT_Float32)
            writeArray2Tiff(lons, inRes, inUL, ls.proj4, lon_fName, gdal.GDT_Float32)
        g = gdal.Open(lat_fName, GA_ReadOnly)
        lat = g.ReadAsArray(xStart, yStart, xSize, ySize)
        g = None

        g = gdal.Open(lon_fName, GA_ReadOnly)
        lon = g.ReadAsArray(xStart, yStart, xSize, ySize)
        g = None

        # ====get overpass hour insolation======================================
        sceneDir = os.path.join(self.satscene_path, 'INSOL')
        outFN = os.path.join(sceneDir, '%s_Insol1.tiff' % sceneID)
        g = gdal.Open(outFN, GA_ReadOnly)
        Rs_1 = g.ReadAsArray(xStart, yStart, xSize, ySize) * 0.042727217
        g = None

        # ====get daily insolation=========================================
        outFN = os.path.join(sceneDir, '%s_Insol24.tiff' % sceneID)
        g = gdal.Open(outFN, GA_ReadOnly)
        Rs24 = g.ReadAsArray(xStart, yStart, xSize, ySize)
        g = None

        # ===============get biophysical parameters at overpass time============
        sceneDir = os.path.join(self.satscene_path, 'ALBEDO')
        #        sceneDir = os.path.join(self.landsatDataBase,'albedo',scene)
        outFN = os.path.join(sceneDir, '%s_albedo.tiff' % sceneID)
        g = gdal.Open(outFN, GA_ReadOnly)
        albedo = g.ReadAsArray(xStart, yStart, xSize, ySize)
        g = None

        # ------>get LAI...
        #        outFN = os.path.join(self.landsatDataBase,'LAI',scene,'lndlai.%s.hdf' % sceneID)
        sceneDir = os.path.join(self.satscene_path, 'LAI')
        outFN = os.path.join(sceneDir, '%s_lai.tiff' % sceneID)
        g = gdal.Open(outFN, GA_ReadOnly)
        LAI = g.ReadAsArray(xStart, yStart, xSize, ySize) * 0.001  # TESTING
        g = None
        #        LAI[np.where(LAI==-9.999)]=np.nan
        #        LAI[np.where(LAI<=0.)]=0.001

        #        hdf = SD(outFN,SDC.READ)
        #        data2D = hdf.select('LAI')
        #        LAI = data2D[yStart:yStart+ySize,xStart:xStart+xSize].astype(np.double)*0.001
        #
        #        LAI[np.where(LAI==-9.999)]=np.nan
        #        LAI[np.where(LAI<=0.)]=0.001

        # ------>get ndvi...'
        #        sceneDir = os.path.join(self.landsatNDVI,scene)
        sceneDir = os.path.join(self.satscene_path, 'NDVI')
        outFN = os.path.join(sceneDir, '%s_ndvi.tiff' % sceneID)
        g = gdal.Open(outFN, GA_ReadOnly)
        ndvi = g.ReadAsArray(xStart, yStart, xSize, ySize)  # *0.001 # TESTING
        g = None
        #        ndvi[np.where(ndvi==-9999.)]=np.nan

        #        data2D = hdf.select('NDVI')
        #        ndvi = data2D[yStart:yStart+ySize,xStart:xStart+xSize].astype(np.double)*0.001
        #        ndvi[np.where(ndvi==-9.999)]=np.nan

        # ===get cfmask=======
        sceneDir = os.path.join(self.satscene_path, 'CF_MASK')
        outFN = os.path.join(sceneDir, '%s_Mask.tiff' % sceneID)
        g = gdal.Open(outFN, GA_ReadOnly)
        cfmask = g.ReadAsArray(xStart, yStart, xSize, ySize)
        g = None
        #        laiFN = os.path.join(self.landsatDataBase,'LAI',scene,'lndlai.%s.hdf' % sceneID)
        #        hdf = SD(laiFN,SDC.READ)
        #        data2D = hdf.select('cfmask')
        #        cfmask = data2D[yStart:yStart+ySize,xStart:xStart+xSize].astype(np.double)

        # ---------->get LST...
        sceneDir = os.path.join(self.satscene_path, 'LST')
        outFN = os.path.join(sceneDir, '%s_lstSharp.tiff' % sceneID)
        g = gdal.Open(outFN, GA_ReadOnly)
        # *NOTE: version 0.2.0 forward------>
        # convert from scaled celcius to kelvin int16->float32
        Tr_K = (g.ReadAsArray(xStart, yStart, xSize, ySize) / 100.) + 273.15
        #        Tr_K = g.ReadAsArray(xStart,yStart,xSize,ySize)+273.16 # TESTING
        g = None
        Tr_K[np.where(albedo < 0)] = np.nan
        # ---------->get LC...
        sceneDir = os.path.join(self.satscene_path, 'LC')
        #        sceneDir = os.path.join(self.landsatDataBase,'LC',scene)
        outFN = os.path.join(sceneDir, '%s_LC.tiff' % sceneID)
        g = gdal.Open(outFN, GA_ReadOnly)
        LCdata = g.ReadAsArray(xStart, yStart, xSize, ySize)
        g = None
        # ---------->get ALEXI mask...
        #        ET_ALEXI[np.where(albedo<0)]=-9999
        #        mask = ET_ALEXI.copy()
        mask = np.tile(1., albedo.shape)
        mask[cfmask > 0] = 0.
        mask[albedo < 0] = 0.
        #        mask[mask==0]=1.
        #        mask[mask==-9999.] = 0.
        albedo[np.where(albedo < 0.)] = np.nan

        # ====================get LC based variables===============================
        s = ndimage.__file__
        envPath = os.sep.join(s.split(os.sep)[:-6])
        landsatLC = os.path.join(envPath, 'share', 'disalexi')
        aleafv = getParFromExcel(LCdata, landsatLC, landcover, 'aleafv')
        aleafn = getParFromExcel(LCdata, landsatLC, landcover, 'aleafn')
        aleafl = getParFromExcel(LCdata, landsatLC, landcover, 'aleafl')
        adeadv = getParFromExcel(LCdata, landsatLC, landcover, 'adeadv')
        adeadn = getParFromExcel(LCdata, landsatLC, landcover, 'adeadn')
        adeadl = getParFromExcel(LCdata, landsatLC, landcover, 'adeadl')
        hc_min = getParFromExcel(LCdata, landsatLC, landcover, 'hmin')
        hc_max = getParFromExcel(LCdata, landsatLC, landcover, 'hmax')
        xl = getParFromExcel(LCdata, landsatLC, landcover, 'xl')
        clump = getParFromExcel(LCdata, landsatLC, landcover, 'omega')
        clump[clump == 0] = 0.99
        LAI[LCdata == 11] = 0.01
        ndvi[LCdata == 11] = -0.5

        aleafv[np.isnan(aleafv)] = 0.9
        aleafn[np.isnan(aleafn)] = 0.9
        aleafl[np.isnan(aleafl)] = 0.9
        adeadv[np.isnan(adeadv)] = 0.2
        adeadn[np.isnan(adeadn)] = 0.2
        adeadl[np.isnan(adeadl)] = 0.2
        hc_min[np.isnan(hc_min)] = 0.1
        hc_max[np.isnan(hc_max)] = 0.5
        xl[np.isnan(xl)] = 0.5
        xl[xl == 0.] = 0.5

        F = LAI * clump  # LAI for leafs spherical distribution
        f_c = 1 - (np.exp(-0.5 * F))  # fraction cover at nadir (view=0)
        f_c[f_c <= 0.01] = 0.01
        f_c[f_c >= 0.9] = 0.9

        # ************************************************************************
        # Compute Canopy height and Roughness Parameters
        hc = hc_min + ((hc_max - hc_min) * f_c)

        #        LAI[np.where(LAI==0.0)]=0.001
        vza = np.tile(0.0, np.shape(LAI))
        #        Rs24 = Rs24+500. # FOR TESTING ONLY
        #        Rs24 = (Rs24*0.0864)/24.0 # MERRA
        Rs24 = Rs24 * 0.0864  # GSIP

        leaf_width = xl
        alpha_PT = np.tile(1.32, np.shape(LAI))
        time = self.dt.hour + (self.dt.minute / 60.)
        #        print("time:%f" % time)
        t_rise, t_end, zs = sunset_sunrise(self.dt, np.deg2rad(lon), np.deg2rad(lat), time)
        #        print("t_rise:%f, t_end:%f" % (t_rise[0,0],t_end[0,0]))
        #        zs = np.tile(sz,np.shape(LAI))

        # ================RUN DisALEXI=================================

        if TSEB_only == 1:
            # convert TA from scaled celcius to kelvin
            #            T_A_K = (T_A_K/1000.)+273.15  # removed /1000 FOR TESTING!!!!
            #            T_A_K = (T_A_K)+273.16  # removed /1000 FOR TESTING!!!!
            nan_check = np.sum(np.isnan(LAI)) / LAI.size
            if nan_check == 1:  # All nans
                ET_24 = np.tile(np.nan, LAI.shape)
            else:
                output = TSEB_PT_usda(
                    Tr_K,
                    vza,
                    T_A_K,
                    u,
                    p,
                    Rs_1,
                    zs,
                    aleafv,
                    aleafn,
                    aleafl,
                    adeadv,
                    adeadn,
                    adeadl,
                    albedo,
                    ndvi,
                    LAI,
                    clump,
                    hc,
                    mask,
                    time,
                    t_rise,
                    t_end,
                    leaf_width=leaf_width,
                    a_PT_in=alpha_PT)  # atmospheric emissivity (clear-sly) Idso and Jackson (1969)

                scaling = 1.0
                Fsun = (output[4] + output[6]) / Rs_1
                #            Rs24 = ndimage.gaussian_filter(Rs24, sigma=5)
                EFeq = Fsun * (Rs24)
                #            ET_24 = EFeq/2.45*scaling
                ET_24 = EFeq * 0.408 * scaling
                ET_24[ET_24 < 0.01] = 0.01
        #            ET_24 = np.array(ET_24*1000.,dtype='uint16')
        else:
            nan_check = np.sum(np.isnan(LAI)) / LAI.size
            if nan_check == 1:  # All nans
                T_A_K = np.tile(np.nan, LAI.shape)
            else:
                output = self.DisALEXI_PT(
                    ET_ALEXI,
                    Rs_1,
                    Rs24,
                    Tr_K,
                    Ta,
                    vza,
                    u,
                    p,
                    zs,
                    aleafv,
                    aleafn,
                    aleafl,
                    adeadv,
                    adeadn,
                    adeadl,
                    albedo,
                    ndvi,
                    LAI,
                    clump,
                    hc,
                    mask,
                    time,
                    t_rise,
                    t_end,
                    leaf_width=leaf_width,
                    alpha_PT=alpha_PT)
                #            T_A_K= np.array((output['T_A_K']-273.15)*1000.,dtype='uint16')
                T_A_K = np.array(output['T_A_K'], dtype='float32')

        #        outFormat = gdal.GDT_UInt16
        outFormat = gdal.GDT_Float32
        outET24Path = os.path.join(self.resultsBase, scene)
        if not os.path.exists(outET24Path):
            os.makedirs(outET24Path)
        # set ouput location and resolution
        #        ulx = self.meta.CORNER_UL_PROJECTION_X_PRODUCT
        #        uly = self.meta.CORNER_UL_PROJECTION_Y_PRODUCT
        #        delx = self.meta.GRID_CELL_SIZE_REFLECTIVE
        #        dely = self.meta.GRID_CELL_SIZE_REFLECTIVE
        ulx = ls.ulx
        uly = ls.uly
        delx = ls.delx
        dely = -ls.dely
        inUL = [ulx + (xStart * delx), uly - (yStart * dely)]
        inRes = [delx, dely]

        if TSEB_only == 1:
            outFormat = gdal.GDT_Float32  # FOR TESTING WE CAN GO BACK TO INT LATER
            ET_24outName = 'ETd_%s_part_%d_%d.tif' % (yeardoy, xStart, yStart)
            fName = '%s%s%s' % (outET24Path, os.sep, ET_24outName)
            writeArray2Tiff(ET_24, inRes, inUL, ls.proj4, fName, outFormat)

            # ======write out fluxes==================================
            #            flag, Ts, Tc, Tac, lETc, H_c, lEs, H_s, G0
            outFormat = gdal.GDT_Float32
            # ==Ts=====>
            Ts_24outName = 'Ts_%s_part_%d_%d.tif' % (yeardoy, xStart, yStart)
            fName = '%s%s%s' % (outET24Path, os.sep, Ts_24outName)
            writeArray2Tiff(output[1], inRes, inUL, ls.proj4, fName, outFormat)
            # ==Tc=====>
            Tc_24outName = 'Tc_%s_part_%d_%d.tif' % (yeardoy, xStart, yStart)
            fName = '%s%s%s' % (outET24Path, os.sep, Tc_24outName)
            writeArray2Tiff(output[2], inRes, inUL, ls.proj4, fName, outFormat)
            # ==Tac=====>
            Tac_24outName = 'Tac_%s_part_%d_%d.tif' % (yeardoy, xStart, yStart)
            fName = '%s%s%s' % (outET24Path, os.sep, Tac_24outName)
            writeArray2Tiff(output[3], inRes, inUL, ls.proj4, fName, outFormat)
            # ==lETc=====>
            lETc_24outName = 'lETc_%s_part_%d_%d.tif' % (yeardoy, xStart, yStart)
            fName = '%s%s%s' % (outET24Path, os.sep, lETc_24outName)
            writeArray2Tiff(output[4], inRes, inUL, ls.proj4, fName, outFormat)
            # ==H_c=====>
            H_c_24outName = 'H_c_%s_part_%d_%d.tif' % (yeardoy, xStart, yStart)
            fName = '%s%s%s' % (outET24Path, os.sep, H_c_24outName)
            writeArray2Tiff(output[5], inRes, inUL, ls.proj4, fName, outFormat)
            # ==lEs=====>
            lEs_24outName = 'lEs_%s_part_%d_%d.tif' % (yeardoy, xStart, yStart)
            fName = '%s%s%s' % (outET24Path, os.sep, lEs_24outName)
            writeArray2Tiff(output[6], inRes, inUL, ls.proj4, fName, outFormat)
            # ==H_s=====>
            H_s_24outName = 'H_s_%s_part_%d_%d.tif' % (yeardoy, xStart, yStart)
            fName = '%s%s%s' % (outET24Path, os.sep, H_s_24outName)
            writeArray2Tiff(output[7], inRes, inUL, ls.proj4, fName, outFormat)
            # ==G0=====>
            G0_24outName = 'G0_%s_part_%d_%d.tif' % (yeardoy, xStart, yStart)
            fName = '%s%s%s' % (outET24Path, os.sep, G0_24outName)
            writeArray2Tiff(output[8], inRes, inUL, ls.proj4, fName, outFormat)

        else:
            T_A_KoutName = 'Ta_%s_%d_%d.tif' % (yeardoy, xStart, yStart)
            fName = '%s%s%s' % (outET24Path, os.sep, T_A_KoutName)
            writeArray2Tiff(T_A_K, inRes, inUL, ls.proj4, fName, outFormat)


#!/usr/bin/env python

"""
Module containing analysis classes which compute a pourbaix diagram given a
target compound/element.
"""

from __future__ import division

__author__ = "Sai Jayaraman"
__copyright__ = "Copyright 2012, The Materials Project"
__version__ = "0.0"
__maintainer__ = "Sai Jayaraman"
__email__ = "sjayaram@mit.edu"
__status__ = "Development"
__date__ = "Nov 1, 2012"

import logging
import numpy as np
import itertools
import re
from itertools import chain
from pyhull.convex_hull import ConvexHull
from pymatgen.analysis.pourbaix.entry import MultiEntry
from pymatgen.core.periodic_table import Element
from pymatgen.core import Composition
from pymatgen.core.ion import Ion

logger = logging.getLogger(__name__)

PREFAC = 0.0591
MU_H2O = -2.4583

class PourbaixDiagram(object):
    """
    Class to create a Pourbaix diagram from entries
    """
    def __init__(self, entries, comp_dict=None):
        """
        Args:
            entries:
                Entries list containing both Solids and Ions
            comp_dict:
                Dictionary of compositions
        """
        self._solid_entries = list()
        self._ion_entries = list()
        for entry in entries:
            if entry.phase_type == "Solid":
                self._solid_entries.append(entry)
            elif entry.phase_type == "Ion":
                self._ion_entries.append(entry)
            else:
                raise StandardError("Incorrect Phase type - needs to be \
                Pourbaix entry of phase type Ion/Solid")
        self._unprocessed_entries = self._solid_entries + self._ion_entries
        self._elt_comp = comp_dict
        if comp_dict:
            self._multielement = True
            self.pourbaix_elements = [key for key in comp_dict]
            w = [comp_dict[key] for key in comp_dict]
            A = []
            for comp in comp_dict:
                m = re.search(r"\[([^\[\]]+)\]|\(aq\)", comp)
                if m:
                    comp_obj = Ion.from_formula(comp)
                else:
                    comp_obj = Composition.from_formula(comp)
                Ai = []
                for elt in self.pourbaix_elements:
                    Ai.append(comp_obj[Element(elt)])
                A.append(Ai)
            A = np.array(A).T.astype(float)
            w = np.array(w)
            A /= np.dot([A[i].sum() for i in xrange(len(A))], w)
            x = np.linalg.solve(A, w)
            self._elt_comp = dict(zip(self.pourbaix_elements, x))

        else:
            self._multielement = False
            self.pourbaix_elements = [el.symbol
                                      for el in entries[0].composition.elements
                                      if el.symbol not in ["H", "O"]]
        self._make_pourbaixdiagram()

    def _create_conv_hull_data(self):
        """
        Make data conducive to convex hull generator.
        """
        if self._multielement:
            self._all_entries = self._process_multielement_entries()
        else:
            self._all_entries = self._unprocessed_entries
        entries_to_process = list()
        for entry in self._all_entries:
            entry.scale(entry.normalization_factor)
            entry.correction += (- MU_H2O * entry.nH2O + entry.conc_term)
            entries_to_process.append(entry)
        self._qhull_entries = entries_to_process
        return self._process_conv_hull_data(entries_to_process)

    def _process_conv_hull_data(self, entries_to_process):
        """
        From a sequence of ion+solid entries, generate the necessary data
        for generation of the convex hull.
        """
        data = []
        for entry in entries_to_process:
            row = [entry.npH, entry.nPhi, entry.g0]
            data.append(row)
        temp = zip(data, self._qhull_entries)
        temp.sort(key=lambda x: x[0][2])
        [data, self._qhull_entries] = zip(*temp)
        return data

    def _process_multielement_entries(self):
        """
        Create entries for multi-element Pourbaix construction
        """
        N = len(self._elt_comp)  # No. of elements
        entries = self._unprocessed_entries
        el_list = self._elt_comp.keys()
        comp_list = [self._elt_comp[el] for el in el_list]
        list_of_entries = list(itertools.combinations(
            [i for i in xrange(len(entries))], N))
        processed_entries = list()
        self._entry_components_list = list_of_entries
        self._entry_components_dict = {}
        count = 0
        for entry_list in list_of_entries:
            # Check if all elements in composition list are present in entry_list
            if not (set([Element(el) for el in el_list]).issubset(set(list(chain.from_iterable([entries[i].composition.keys() for i in entry_list]))))):
                continue
            count += 1
            A = [[0.0] * (len(el_list) - 1) for _ in range(len(entry_list) - 1)]
            multi_entries = [entries[j] for j in entry_list]
            entry0 = entries[entry_list[0]]
            comp0 = entry0.composition
            if entry0.phase_type == "Solid":
                red_fac = comp0.get_reduced_composition_and_factor()[1]
            else:
                red_fac = 1.0
            sum_nel = sum([comp0[el] / red_fac for el in el_list])
            b = [comp0[Element(el_list[i])] / red_fac - comp_list[i] * sum_nel
                 for i in xrange(1, len(el_list))]
            for j in xrange(1, len(entry_list)):
                entry = entries[entry_list[j]]
                comp = entry.composition
                if entry.phase_type == "Solid":
                    red_fac = comp.get_reduced_composition_and_factor()[1]
                else:
                    red_fac = 1.0
                sum_nel = sum([comp[el] / red_fac for el in el_list])
                for i in xrange(1, len(el_list)):
                    el = el_list[i]
                    A[i-1][j-1] = comp_list[i] * sum_nel -\
                        comp[Element(el)] / red_fac
            try:
                weights = np.linalg.solve(np.array(A), np.array(b))
            except np.linalg.linalg.LinAlgError as err:
                if 'Singular matrix' in err.message:
                    continue
                else:
                    raise StandardError("Unknown Error message!")
            if not(np.all(weights > 0.0)):
                continue
            weights = list(weights)
            weights.insert(0, 1.0)
            super_entry = MultiEntry(multi_entries, weights)
            self._entry_components_dict[super_entry] = entry_list
            processed_entries.append(super_entry)
        return processed_entries

    def _make_pourbaixdiagram(self):
        """
        Calculates entries on the convex hull in the dual space.
        """
        stable_entries = set()
        self._qhull_data = self._create_conv_hull_data()
        dim = len(self._qhull_data[0])
        if len(self._qhull_data) < dim:
            raise StandardError("Can only do elements with at-least 3 entries"
                                " for now")
        if len(self._qhull_data) == dim:
            self._facets = [range(dim)]
        else:
            facets_pyhull = np.array(ConvexHull(self._qhull_data).vertices)
            self._facets = np.sort(np.array(facets_pyhull))
            logger.debug("Final facets are\n{}".format(self._facets))

            logger.debug("Removing vertical facets...")
            vert_facets_removed = list()
            for facet in self._facets:
                facetmatrix = np.zeros((len(facet), len(facet)))
                count = 0
                for vertex in facet:
                    facetmatrix[count] = np.array(self._qhull_data[vertex])
                    facetmatrix[count, dim - 1] = 1
                    count += 1
                if abs(np.linalg.det(facetmatrix)) > 1e-8:
                    vert_facets_removed.append(facet)
                else:
                    print "removed facet", facet
                    logger.debug("Removing vertical facet : {}".format(facet))

            logger.debug("Removing UCH facets by eliminating normal.z >0 ...")

            # Find center of hull
            vertices = set()
            for facet in vert_facets_removed:
                for vertex in facet:
                    vertices.add(vertex)
            c = [0.0, 0.0, 0.0]
            c[0] = np.average([self._qhull_data[vertex][0]
                               for vertex in vertices])
            c[1] = np.average([self._qhull_data[vertex][1]
                               for vertex in vertices])
            c[2] = np.average([self._qhull_data[vertex][2]
                               for vertex in vertices])

            # Shift origin to c
            new_qhull_data = np.array(self._qhull_data)
            for vertex in vertices:
                new_qhull_data[vertex] -= c

            # For each facet, find normal n, find dot product with P, and
            # check if this is -ve
            final_facets = list()
            for facet in vert_facets_removed:
                a = new_qhull_data[facet[1]] - new_qhull_data[facet[0]]
                b = new_qhull_data[facet[2]] - new_qhull_data[facet[0]]
                n = np.cross(a, b)
                val = np.dot(n, new_qhull_data[facet[0]])
                if val < 0:
                    n = -n
                if n[2] <= 0:
                    final_facets.append(facet)
                else:
                    print "removed UCH facet", facet
                    logger.debug("Removing UCH facet : {}".format(facet))
            final_facets = np.array(final_facets)
            self._facets = final_facets

        stable_vertices = set()
        for facet in self._facets:
            for vertex in facet:
                stable_vertices.add(vertex)
                stable_entries.add(self._qhull_entries[vertex])
        self._stable_entries = stable_entries
        self._vertices = stable_vertices

    @property
    def facets(self):
        """
        Facets of the convex hull in the form of  [[1,2,3],[4,5,6]...]
        """
        return self._facets

    @property
    def qhull_data(self):
        """
        Data used in the convex hull operation. This is essentially a matrix of
        composition data and energy per atom values created from qhull_entries.
        """
        return self._qhull_data

    @property
    def qhull_entries(self):
        """
        Return qhull entries
        """
        return self._qhull_entries

    @property
    def stable_entries(self):
        """
        Returns the stable entries in the phase diagram.
        """
        return self._stable_entries

    @property
    def all_entries(self):
        """
        Return all entries
        """
        return self._all_entries

    @property
    def vertices(self):
        """
        Return vertices of the convex hull
        """
        return self._vertices

    @property
    def unprocessed_entries(self):
        """
        Return unprocessed entries
        """
        return self._unprocessed_entries


#!/usr/bin/env python
import argparse
from pathlib import Path
from astropy.time import Time
from scipy.optimize import leastsq
import fitsio
import numpy as np
import textwrap
from . import sdss_paths

try:
    from bin import epics_fetch
except ImportError as e:
    raise ImportError('Please add ObserverTools/bin to your PYTHONPATH:\n'
                      '    {}'.format(e))

__version__ = '3.2.1'

class APOGEERaw:
    """A class to parse raw data from APOGEE. The purpose of this class is to
    read raw image files from /data/apogee/archive, regardless of any future
    changes in SDSS. This
    will hopefully help SDSS-V logging.

    Methods:
        compute_offets: Returns the offsets of an arc/object to compute the
        relative dither

        ap_test: Returns a list of faint fibers and missing fibers from a flat
        """

    def __init__(self, fil, args, ext=1, ):
        self.file = Path(fil)
        if not self.file.exists():
            raise FileNotFoundError(f"Could not find file {self.file.absolute()}")
        self.ext = ext
        self.args = args
        header = fitsio.read_header(fil, ext=ext)
        self.telemetry = epics_fetch.telemetry
        dithers = self.telemetry.get('25m:apogee:ditherNamedPositions',
                                     start=(Time.now() - 5 / 24 / 60).datetime,
                                     end=Time.now().datetime,
                                     scan_archives=False, interpolation='raw')
        # layer = self.image[layer_ind]
        # An A dither is DITHPIX=12.994, a B dither is DITHPIX=13.499
        if (header['DITHPIX'] - dithers.values[-1][0]) < 0.05:
            self.dither = 'A'
        elif (header['DITHPIX'] - dithers.values[-1][1]) < 0.05:
            self.dither = 'B'
        else:
            self.dither = '{:.1f}'.format(header['DITHPIX'])
        self.exp_time = header['EXPTIME']
        self.isot = Time(header['DATE-OBS'])  # Local
        self.plate_id = header['PLATEID']
        self.cart_id = header['CARTID']
        self.exp_id = int(str(fil).split('-')[-1].split('.')[0])
        if header['EXPTYPE'].capitalize() == 'Arclamp':
            if header['LAMPUNE']:
                self.exp_type = 'UNe Arc'
            elif header['LAMPTHAR']:
                self.exp_type = 'ThAr Arc'
            elif "FPI" in header["OBSCMNT"]:
                self.exp_type = "FPI Lamp"
            else:
                print('Could not process exposure type of {}'.format(self.file))
        else:
            self.exp_type = header['EXPTYPE'].capitalize()
        self.n_read = header['NREAD']
        self.lead = header['PLATETYP']
        self.img_type = header['IMAGETYP'].capitalize()
        if header['EXPTYPE'] == 'OBJECT':
            self.seeing = header['SEEING']
        else:
            self.seeing = 0.0

        self.quickred_data = np.array([[]])
        self.quickred_file = ''
        self.utr_file = ''
        self.utr_data = np.array([[]])

    # noinspection PyTupleAssignmentBalance,PyTypeChecker
    def compute_offset(self, fibers=(30, 35), w0=939, dw=40, sigma=1.2745):
        """This is based off of apogeeThar.OneFileFitting written by Elena. It
        is supposed to generate a float for the pixel offsets of an APOGEE
        ThAr cal. Here is how it works:
        It opens a quickred file, which is of shape n_fiber*n_dispersion_pixels,
        and then it averages the fibers inside the fibers tuple. It then based
        off of the provided w0 (mean) and sigma, it creates a gaussian function
        and then creates a function called err_func that compares the gaussian
        with a slice of the data, from w0-dw/2 to w0+dw/2. It then uses scipy's
        least squared equation solver to find the difference between w0 given as
        an input and the actual w0 of the spectral line. This only works if you
        pick a prominent line to go off of. The default parameters are given for
        ThAr lines, but UNe lines could also be used, with the following inputs:
        fibers: (30, 35)
        w0: 1761
        dw: 20
        sigma: 3
        """
        w0 = int(w0)
        dw = int(dw)
        mjd = self.file.absolute().parent.name
        self.quickred_file = (self.file.absolute().parent.parent.parent
                              / 'quickred/{}/ap1D-a-{}.fits.fz'
                                ''.format(mjd, self.exp_id))
        try:
            if not self.quickred_data:
                self.quickred_data = fitsio.read(self.quickred_file, 1)
        except OSError as e:
            if self.args.verbose:
                print('Offsets for {} produced this error\n{}'.format(self.file,
                                                                      e))
            return np.nan
        lower = w0 - dw // 2
        upper = w0 + dw // 2
        line_inds = np.arange(self.quickred_data.shape[1])[lower:upper]
        line = np.average(self.quickred_data[fibers[0]:fibers[1], lower:upper],
                          axis=0)

        def fit_func(w, x):
            return np.exp(-0.5 * ((x - w) / sigma) ** 2)

        def err_func(w, x, y):
            return fit_func(w, x) - y

        w_model, success = leastsq(err_func, w0, args=(line_inds, line))

        diff = w_model[0] - w0
        if np.abs(diff) > 10:
            print('A large dither was reported for exposure {}: {:.3f}'
                  '\n  This may mean the zero point needs to be'
                  ' adjusted, currently it is {}'
                  ''.format(self.exp_id, diff, w0))
        return diff

    def ap_test(self, ws=(900, 910), master_col=None, plot=False, legacy=False,
                dome_flat_shape=None, n_fibers=300, print_it=False):
        if master_col is None:
            raise ValueError("APTest didn't receive a valid master_col: {}"
                             "".format(master_col))
        if legacy:
            mjd = self.file.absolute().parent.name
            self.utr_file = sdss_paths.ap_utr / f"{mjd}/apRaw-{self.exp_id}.fits"
            if not self.utr_file.exists():
                raise FileNotFoundError(f"Couldn't fine the file: {self.utr_file.as_posix()}")
            try:
                self.utr_data = fitsio.read(self.utr_file, 0)
            except OSError as e:
                if self.args.verbose:
                    print('APTest for {} produced this error\n{}'.format(
                        self.file, e))
            slc0 = np.average(self.utr_data[::-1, ws[0]:ws[1]], axis=1)
            # print(slc0.mean(), slc0.shape, slc0[0])
            slc = np.zeros(n_fibers)
            for j in range(n_fibers):
                for k in range(10):
                    if dome_flat_shape[j, k] != 0:
                        slc[j] += slc0[dome_flat_shape[j, k]]
            # print(slc.mean(), slc.shape, slc[100])

        else:
            if self.quickred_data.size == 0:
                mjd = self.file.absolute().parent.name
                self.quickred_file = (self.file.absolute().parent.parent.parent
                                      / 'quickred/{}/ap1D-a-{}.fits.fz'
                                        ''.format(mjd, self.exp_id))
                try:
                    self.quickred_data = fitsio.read(self.quickred_file, 1)
                except OSError as e:
                    if self.args.verbose:
                        print('APTest for {} produced this error\n{}'.format(
                            self.file, e))
            slc = np.average(self.quickred_data[:, ws[0]:ws[1]], axis=1)
        # print(master_col.mean(), master_col.shape, master_col[100])
        flux_ratio = slc / master_col
        # flux_ratio = flux_ratio / flux_ratio.sum() / flux_ratio.shape[0]
        flux_ratio = flux_ratio / flux_ratio.sum() * flux_ratio.shape[0]
        # print(flux_ratio.mean(), flux_ratio.shape, flux_ratio[100])
        bad_data = ((flux_ratio == np.inf)
                    | (flux_ratio == -np.inf)
                    | np.isnan(flux_ratio))
        flux_ratio[bad_data] = np.nan
        avg = np.nanmean(flux_ratio)
        missing = flux_ratio < 0.2
        faint = (flux_ratio < 0.7) & (0.2 <= flux_ratio)
        bright = ~missing & ~faint
        i_missing = np.where(missing)[0].astype(int) + 1
        i_faint = np.where(faint)[0].astype(int) + 1
        i_bright = np.where(bright)[0]
        missing_bundles = self.create_bundles(i_missing)
        faint_bundles = self.create_bundles(i_faint)
        if print_it:

            print(textwrap.fill('Missing Fibers: {}'.format(missing_bundles),
                                80))
            print(textwrap.fill('Faint Fibers: {}'.format(faint_bundles), 80))
            print()

        if plot:
            import matplotlib.pyplot as plt
            fig = plt.figure(figsize=(9, 4))
            ax = fig.gca()
            x = np.arange(len(flux_ratio)) + 1
            ax.plot(x[i_bright], flux_ratio[i_bright], 'o', c=(0, 0.6, 0.533))
            ax.plot(x[i_faint], flux_ratio[i_faint], 'o', c=(0.933, 0.466, 0.2))
            ax.plot(x[i_missing], flux_ratio[i_missing], 'o',
                    c=(0.8, 0.2, 0.066))
            ax.set_xlabel('Fiber ID')
            ax.set_ylabel('Throughput Efficiency')
            ax.axis([1, 300, -0.2, 1.35])
            ax.grid(True)
            ax.axhline(0.7, c=(0, 0.6, 0.533))
            ax.axhline(0.2, c=(0.933, 0.466, 0.2))
            ax.set_title('APOGEE Fiber Relative Intensity {}'.format(
                self.exp_id), size=15)
            fig.show()

        return missing_bundles, faint_bundles, avg

    @staticmethod
    def create_bundles(subset):
        """This method converts an array of ints into a list of strings that
        describe a large series of fibers and ints for lone fibers.

        Ex: [1, 2, 3, 5] -> ['1 - 3', 5]

        """
        if len(subset) == 0:
            return []
        bundles = [subset[0]]
        b = 0
        for fib in subset[1:]:
            if isinstance(bundles[b], np.int64):
                if bundles[b] + 1 == fib:
                    if fib % 30 == 0:
                        bundles.append(fib)
                        b += 1
                    else:
                        # All strings are created here
                        bundles[b] = '{} - {}'.format(bundles[b], fib)
                else:
                    bundles.append(fib)
                    b += 1
            elif isinstance(bundles[b], str):
                if int(bundles[b].split()[-1]) + 1 == fib:
                    if fib % 30 == 1:
                        bundles.append(fib)
                        b += 1
                    else:
                        bundles[b] = '{} - {}'.format(
                            bundles[b].split()[0], fib)
                else:
                    bundles.append(fib)
                    b += 1
        # for i, bundle in enumerate(bundles):
        #     if isinstance(bundle, str):
        #         low, high = np.array(bundle.split(' - ')).astype(int)
        #         if ((low - 1) // 30) == ((high - 1) // 30):
        #             bundles[i] = '{} bundle'.format(low)
        return bundles

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-t', '--today', action='store_true',
                        help="Whether or not you want to search for today's"
                             " data, whether or not the night is complete."
                             " Note: must be run after 00:00Z")
    parser.add_argument('-m', '--mjd',
                        help='If not today (-t), the mjd to search')
    parser.add_argument('-v', '--verbose', action='count', default=1,
                        help='Show details, can be stacked')
    args = parser.parse_args()
    if args.today:
        mjd_today = int(Time.now().sjd)
        data_dir = sdss_paths.ap_archive / f"{mjd_today}/"
    elif args.mjd:
        data_dir = sdss_paths.ap_archive / f"{args.mjd}"
    else:
        raise Exception('No date specified')
    # print(data_dir)
    for path in data_dir.rglob('apR*.apz'):
        print(path)

if __name__ == '__main__':
    main()


"""
model the Subaru optics system

This is a code modified from Rupert's original optics_propagate.py. This code adds more optics to the system,
as well as puts the AO, etc in order for Subaru.

Here, we will add the basic functionality of the Subaru Telescope, including the primary, secondary, and AO188.
The SCExAO system sits behind the AO188 instrument of Subaru, which is a 188-element AO system located at the
Nasmyth focus (IR) of the telescope. AO188 is a basic 4f-type optical system with a 188 element DM (not perfectly
represented here. Proper will only simulate 2D square DM's, so we use a 14x14 square) inserted in the middle of the
collimated beam. This routine simulates an idealized detector at the A0188 focus.

If you want to scale the intensity of the companion, you can't then normalize the intensity
(prop_define_entrance) since the normalization is per wavefront, so the companion will normalize to itself
rather than to the star. Doing it this way may cause confusion, but better than re-writing things in proper

AO188 uses laser guide-star technology. More info here:
https://subarutelescope.org/Introduction/instrument/AO188.html

A more detailed model of SCExAO will be modelled in a SCExAO_optics.py code. However, this routine is designed for
simple simulations that need to optimize runtime but still have relevance to the Subaru Telescope.

This script is meant to override any Subaru/SCExAO-specific parameters specified in the user's params.py
"""

import numpy as np
from inspect import getframeinfo, stack
import proper

from medis.params import ap, tp, sp
from medis.utils import dprint
import medis.optics as opx
import medis.aberrations as aber
import medis.adaptive as ao
import medis.atmosphere as atmos

#################################################################################################
#################################################################################################
#################################################################################################
# Defining Subaru parameters
# ----------------------------
# According to Iye-et.al.2004-Optical_Performance_of_Subaru:AstronSocJapan, the AO188 uses the IR-Cass secondary,
# but then feeds it to the IR Nasmyth f/13.6 focusing arrangement. So instead of simulating the full Subaru system,
# we can use the effective focal length at the Nasmyth focus, and simulate it as a single lens.
tp.d_nsmyth = 7.9716  # m pupil diameter
tp.fn_nsmyth = 13.612  # f# Nasmyth focus
tp.flen_nsmyth = tp.d_nsmyth * tp.fn_nsmyth  # m focal length
tp.dist_nsmyth_ao1 = tp.flen_nsmyth + 1.14  # m distance secondary to M1 of AO188 (hand-tuned, could update with
                                            # data from literature)

#  Below are the actual dimenstions of the Subaru telescope.
# --------------------------------
# tp.enterence_d = 8.2  # m diameter of primary
# tp.flen_primary = 15  # m focal length of primary
# tp.dist_pri_second = 12.652  # m distance primary -> secondary
# Secondary
tp.d_secondary = 1.265  # m diameter secondary, used for central obscuration
# tp.fn_secondary = 12.6

# Re-writing params terms in Subaru-units
# need this to accurately make atmospheric and aberration maps
tp.entrance_d = tp.d_nsmyth
tp.flen_primary = tp.flen_nsmyth

# Effective Primary Aberrations
# primary_aber_vals = {'a': [7.2e-17, 3e-17],  # power at low spatial frequencies (m4)
#                      'b': [0.8, 0.2],  # correlation length (b/2pi defines knee)
#                      'c': [3.1, 0.5],  #
#                      'a_amp': [0.05, 0.01]}
# ----------------------------
# AO188 OAP1
# Paramaters taken from "Design of the Subaru laser guide star adaptive optics module"
#  Makoto Watanabe et. al. SPIE doi: 10.1117/12.551032
tp.d_ao1 = 0.20  # m  diamater of AO1
tp.fl_ao1 = 1.201  # m  focal length OAP1
tp.dist_ao1_dm = 1.345  # m distance OAP1 to DM
# OAP1_aber_vals = {'a': [7.2e-17, 3e-17],  # power at low spatial frequencies (m4)
#                   'b': [0.8, 0.2],  # correlation length (b/2pi defines knee)
#                   'c': [3.1, 0.5],  #
#                   'a_amp': [0.05, 0.01]}

# ----------------------------
# AO188 OAP2
tp.dist_dm_ao2 = 2.511-tp.dist_ao1_dm  # m distance DM to OAP2
tp.d_ao2 = 0.2  # m  diamater of AO2
tp.fl_ao2 = 1.201  # m  focal length AO2
tp.dist_oap2_focus = 1.261
# OAP2_aber_vals = {'a': [7.2e-17, 3e-17],  # power at low spatial frequencies (m4)
#                   'b': [0.8, 0.2],  # correlation length (b/2pi defines knee)
#                   'c': [3.1, 0.5],  #
#                   'a_amp': [0.05, 0.01]}

tp.lens_params = [{'aber_vals': [7.2e-17, 0.8, 3.1],
                   'diam': tp.entrance_d,
                   'fl': tp.flen_nsmyth,
                   'dist': tp.dist_nsmyth_ao1,
                   'name': 'effective-primary'},

                  {'aber_vals': [7.2e-17, 0.8, 3.1],
                   'diam': tp.d_ao1,
                   'fl': tp.fl_ao1,
                   'dist': tp.dist_ao1_dm,
                   'name': 'ao188-OAP1'},

                  {'aber_vals': [7.2e-17, 0.8, 3.1],
                   'diam': tp.d_ao2,
                   'fl': tp.fl_ao2,
                   'dist': tp.dist_oap2_focus,
                   'name': 'ao188-OAP2'}
                    ]
#################################################################################################
#################################################################################################
#################################################################################################

def Subaru_frontend(empty_lamda, grid_size, PASSVALUE):
    """
    propagates instantaneous complex E-field thru Subaru from the primary through the AO188
        AO system in loop over wavelength range

    this function is called a 'prescription' by proper

    uses PyPROPER3 to generate the complex E-field at the source, then propagates it through atmosphere,
        then telescope, to the focal plane
    the AO simulator happens here
    this does not include the observation of the wavefront by the detector
    :returns spectral cube at instantaneous time in the focal_plane()
    """
    # print("Propagating Broadband Wavefront Through Subaru")

    # Initialize the Wavefront in Proper
    wfo = opx.Wavefronts()
    wfo.initialize_proper()

    # Atmosphere
    # atmos has only effect on phase delay, not intensity
    wfo.loop_collection(atmos.add_atmos, PASSVALUE['iter'], plane_name='atmosphere')

    # Defines aperture (baffle-before primary)
    # Obscurations (Secondary and Spiders)
    wfo.loop_collection(opx.add_obscurations, d_primary=tp.d_nsmyth, d_secondary=tp.d_secondary, legs_frac=0.05)
    wfo.loop_collection(proper.prop_circular_aperture,
                           **{'radius': tp.entrance_d / 2})  # clear inside, dark outside
    wfo.loop_collection(proper.prop_define_entrance, plane_name='entrance_pupil')  # normalizes abs intensity

    if ap.companion:
        # Must do this after all calls to prop_define_entrance
        wfo.loop_collection(opx.offset_companion)
        wfo.loop_collection(proper.prop_circular_aperture,
                               **{'radius': tp.entrance_d / 2})  # clear inside, dark outside

    # Test Sampling
    if sp.verbose:
        opx.check_sampling(PASSVALUE['iter'], wfo, "initial", getframeinfo(stack()[0][0]), units='mm')
    # Testing Primary Focus (instead of propagating to focal plane)
    # wfo.loop_collection(opx.prop_pass_lens, tp.flen_nsmyth, tp.flen_nsmyth)  # test only going to prime focus

    ########################################
    # Subaru Propagation
    #######################################
    # Effective Primary
    # CPA from Effective Primary
    wfo.loop_collection(aber.add_aber, step=PASSVALUE['iter'], lens_name='ao188-OAP1')
    # Zernike Aberrations- Low Order
    # wfo.loop_collection(aber.add_zern_ab, tp.zernike_orders, aber.randomize_zern_values(tp.zernike_orders))
    wfo.loop_collection(opx.prop_pass_lens, tp.flen_nsmyth, tp.dist_nsmyth_ao1)

    ########################################
    # AO188 Propagation
    ########################################
    # # AO188-OAP1
    wfo.loop_collection(aber.add_aber, step=PASSVALUE['iter'], lens_name='ao188-OAP1')
    wfo.loop_collection(opx.prop_pass_lens, tp.fl_ao1, tp.dist_ao1_dm)

    # AO System
    if tp.use_ao:
        WFS_map = ao.open_loop_wfs(wfo)
        wfo.loop_collection(ao.deformable_mirror, WFS_map, PASSVALUE['iter'], plane_name='woofer',
                            debug=sp.debug)  # don't use PASSVALUE['WFS_map'] here because open loop
    # ------------------------------------------------
    wfo.loop_collection(proper.prop_propagate, tp.dist_dm_ao2)

    # AO188-OAP2
    wfo.loop_collection(aber.add_aber, step=PASSVALUE['iter'], lens_name='ao188-OAP2')
    # wfo.loop_collection(aber.add_zern_ab, tp.zernike_orders, aber.randomize_zern_values(tp.zernike_orders)/2)
    wfo.loop_collection(opx.prop_pass_lens, tp.fl_ao2, tp.dist_oap2_focus)

    ########################################
    # Focal Plane
    # #######################################
    # Check Sampling in focal plane
    if sp.verbose:
        wfo.loop_collection(opx.check_sampling, PASSVALUE['iter'], "focal plane",
                            getframeinfo(stack()[0][0]), units='nm')

    # wfo.focal_plane fft-shifts wfo from Fourier Space (origin==lower left corner) to object space (origin==center)
    cpx_planes, sampling = wfo.focal_plane()

    print(f"Finished datacube at timestep = {PASSVALUE['iter']}")

    return cpx_planes, sampling


from __future__ import print_function
import sys
import numpy as np #, scipy as sp, os, gc
from copy import deepcopy
#from warnings import warn
from time import time

class BoundaryCondition(object):
    """Base class for applying all types of boundary conditions"""

    def __init__(self,
        surface_identification_algorithm='minimisation',
        modify_linear_mesh_on_projection=False,
        project_on_curves=True,
        activate_bounding_box=False,
        bounding_box_padding=1e-3,
        has_planar_surfaces=True,
        solve_for_planar_faces=True,
        save_dirichlet_data=False,
        save_nurbs_data=False,
        filename=None,
        read_dirichlet_from_file=False,
        make_loading="ramp",
        compound_dirichlet_bcs=False
        ):

        # TYPE OF BOUNDARY: straight or nurbs
        self.boundary_type = 'straight'
        self.dirichlet_data_applied_at = 'node' # or 'faces'
        self.neumann_data_applied_at = 'node' # or 'faces'
        self.requires_cad = False
        self.cad_file = None
        # PROJECTION TYPE FOR CAD EITHER orthogonal OR arc_length
        self.projection_type = 'orthogonal'
        # WHAT TYPE OF ARC LENGTH BASED PROJECTION, EITHER 'equal' OR 'fekete'
        self.nodal_spacing_for_cad = 'equal'
        self.project_on_curves = project_on_curves
        self.scale_mesh_on_projection = False
        self.scale_value_on_projection = 1.0
        self.condition_for_projection = 1.0e20
        self.has_planar_surfaces = False
        self.solve_for_planar_faces = solve_for_planar_faces
        self.projection_flags = None
        # FIX DEGREES OF FREEDOM EVERY WHERE CAD PROJECTION IS NOT APPLIED
        self.fix_dof_elsewhere = True
        # FOR 3D ARC-LENGTH PROJECTION
        self.orthogonal_fallback_tolerance = 1.0
        # WHICH ALGORITHM TO USE FOR SURFACE IDENTIFICATION, EITHER 'minimisation' or 'pure_projection'
        self.surface_identification_algorithm = surface_identification_algorithm
        # MODIFY LINEAR MESH ON PROJECTION
        self.modify_linear_mesh_on_projection = modify_linear_mesh_on_projection
        # COMPUTE A BOUNDING BOX FOR EACH CAD SURFACE
        self.activate_bounding_box = activate_bounding_box
        self.bounding_box_padding = float(bounding_box_padding)

        # FOR IGAKit WRAPPER
        self.nurbs_info = None
        self.nurbs_condition = None

        self.analysis_type = 'static'
        self.analysis_nature = 'linear'

        self.dirichlet_flags = None
        self.applied_dirichlet = None
        self.is_dirichlet_computed = False
        self.columns_out = None
        self.columns_in = None
        self.save_dirichlet_data = save_dirichlet_data
        self.save_nurbs_data = save_nurbs_data
        self.filename = filename
        self.read_dirichlet_from_file = read_dirichlet_from_file

        self.neumann_flags = None
        self.applied_neumann = None
        self.is_applied_neumann_shape_functions_computed = False

        self.pressure_flags = None
        self.applied_pressure = None
        self.pressure_increment = 1.0
        self.spring_flags = None
        self.applied_spring = None
        self.master_faces = None
        self.slave_faces = None
        self.applied_connector = None
        self.connector_flags = None
        self.connector_elements = None
        self.connector_faces = None
        self.is_body_force_shape_functions_computed = False

        self.make_loading = make_loading # "ramp" or "constant"
        self.has_step_wise_dirichlet_loading = False
        self.step_wise_dirichlet_data = None
        self.has_step_wise_neumann_loading = False
        self.step_wise_neumann_data = None

        self.compound_dirichlet_bcs = compound_dirichlet_bcs

        # STORE A COPY OF SELF AT THE START TO RESET TO AT THE END
        self.__save_state__()
        # FOR INTERNAL PURPOSES WHEN WE DO NOT WANT TO REST
        self.do_not_reset = True

    def __save_state__(self):
        self.__initialdict__ = deepcopy(self.__dict__)

    def SetDirichletCriteria(self, func, *args, **kwargs):
        """Applies user defined Dirichlet data to self
        """

        if "apply" in kwargs.keys():
            del kwargs["apply"]
            self.has_step_wise_dirichlet_loading = True
            self.step_wise_dirichlet_data = {'func':func, 'args': args, 'kwargs': kwargs}
            self.dirichlet_flags = func(0, *args, **kwargs)
            return self.dirichlet_flags

        self.dirichlet_flags = func(*args, **kwargs)
        return self.dirichlet_flags

    def SetNeumannCriteria(self, func, *args, **kwargs):
        """Applies user defined Neumann data to self
        """

        if "apply" in kwargs.keys():
            del kwargs["apply"]
            self.has_step_wise_neumann_loading = True
            self.step_wise_neumann_data = {'func':func, 'args': args, 'kwargs': kwargs}
            tups = func(0, *args, **kwargs)
        else:
            tups = func(*args, **kwargs)

        if not isinstance(tups,tuple) and self.neumann_data_applied_at == "node":
            self.neumann_flags = tups
            return self.neumann_flags
        else:
            self.neumann_data_applied_at == "face"
            if len(tups) !=2:
                raise ValueError("User-defined Neumann criterion function {} "
                    "should return one flag and one data array".format(func.__name__))
            self.neumann_flags = tups[0]
            self.applied_neumann = tups[1]
            return tups

    def SetRobinCriteria(self, func, *args, **kwargs):
        """Applies user defined Robin data to self, just working on surfaces
        """

        dics = func(*args, **kwargs)

        if isinstance(dics,dict):
            self.RobinLoadSelector(dics)
        elif isinstance(dics,tuple):
            for idic in range(len(dics)):
                if isinstance(dics[idic],dict):
                    self.RobinLoadSelector(dics[idic])
                else:
                    raise ValueError("User-defined Robin criterion function {} "
                        "should return dictionary or tuple(dict,dict,...)".format(func.__name__))
        else:
            raise ValueError("User-defined Robin criterion function {} "
                "should return dictionary or tuple".format(func.__name__))

        return dics

    def RobinLoadSelector(self, tups):
        if tups['type'] == 'Pressure':
            self.pressure_flags = tups['flags']
            self.applied_pressure = tups['data']
        elif tups['type'] == 'Spring':
            self.spring_flags = tups['flags']
            self.applied_spring = tups['data']
        elif tups['type'] == 'Connector':
            self.master_faces = tups['master_faces']
            self.slave_faces = tups['slave_faces']
            self.applied_connector = tups['data']
            self.connector_flags = tups['flags']
            if self.master_faces.shape[0] != self.slave_faces.shape[0]:
                raise ValueError("The size of master_faces and slave_faces should be equal")
        elif tups['type'] == 'Dashpot':
            raise ValueError("Surrounding viscoelastic effects not implemented yet")
        else:
            raise ValueError("Type force {} not understood or not available. "
                "Types are Pressure, Spring, SpringJoint and Dashpot.".format(tups['type']))

    def GetConnectorElements(self, mesh):
        """ Receive the faces along the surfaces interacting """

        # gets the points in the dissection surfaces
        master_points = np.unique(mesh.faces[self.master_faces,:])
        slave_points = np.unique(mesh.faces[self.slave_faces,:])
        # array with the coordinate of the master and slave points
        master_points_coor = mesh.points[master_points]
        slave_points_coor = mesh.points[slave_points]
        # look for a connection between master and slave points
        from scipy.spatial import cKDTree
        tree = cKDTree(master_points_coor)
        distance, id_point = tree.query(slave_points_coor,k=1)
        pair_node_master_slave = np.c_[master_points[id_point],slave_points]
        # build the elements
        nodeperface = mesh.faces.shape[1]
        connector_elements = np.zeros((self.master_faces.shape[0],2*nodeperface),np.uint64)
        connector_elements[:,:4] = mesh.faces[self.master_faces]
        # match the master nodes with its slave within the element
        faces_s = np.zeros(self.master_faces.shape[0],dtype=np.uint64)
        for i in range(self.master_faces.shape[0]):
            iface = self.master_faces[i]
            jnode_array = np.zeros(nodeperface,dtype=np.uint64)
            for j in range(nodeperface):
                inode = mesh.faces[iface,j]
                idx = np.where(pair_node_master_slave[:,0]==inode)[0]
                jnode = pair_node_master_slave[idx,1]
                connector_elements[i,j+nodeperface] = jnode
                jnode_array[j] = jnode
            # use the slave point to recover the slave face respect a master face
            jface_array = np.where(mesh.faces==jnode_array[0])[0]
            for k in range(1,jnode_array.shape[0]):
                jface_array = np.append(jface_array, np.where(mesh.faces==jnode_array[k])[0])
            values, counts = np.unique(jface_array,return_counts=True)
            jface = values[np.where(counts==nodeperface)[0]]
            faces_s[i] = jface

        pair_face_master_slave = np.c_[self.master_faces,faces_s]
        pair_face_master_slave = np.array(pair_face_master_slave, dtype=np.uint64, copy=True)

        self.connector_elements = connector_elements
        self.connector_faces = pair_face_master_slave

        return

    def GetDirichletBoundaryConditions(self, formulation, mesh, materials=None, solver=None, fem_solver=None):

        nvar = formulation.nvar
        ndim = formulation.ndim
        self.columns_in, self.applied_dirichlet = [], []

        #----------------------------------------------------------------------------------------------------#
        #-------------------------------------- NURBS BASED SOLUTION ----------------------------------------#
        #----------------------------------------------------------------------------------------------------#
        if self.boundary_type == 'nurbs':

            tCAD = time()

            if self.read_dirichlet_from_file is False:

                if not self.is_dirichlet_computed:
                    # GET DIRICHLET BOUNDARY CONDITIONS BASED ON THE EXACT GEOMETRY FROM CAD
                    if self.requires_cad:
                        # CALL POSTMESH WRAPPER
                        nodesDBC, Dirichlet = self.PostMeshWrapper(formulation, mesh, materials, solver, fem_solver)
                else:
                    nodesDBC, Dirichlet = self.nodesDBC, self.Dirichlet

                # GET DIRICHLET DoFs
                self.columns_out = (np.repeat(nodesDBC,nvar,axis=1)*nvar +\
                 np.tile(np.arange(nvar)[None,:],nodesDBC.shape[0]).reshape(nodesDBC.shape[0],formulation.ndim)).ravel()
                self.applied_dirichlet = Dirichlet.ravel()

                # FIX THE DOF IN THE REST OF THE BOUNDARY
                if self.fix_dof_elsewhere:
                    if ndim==2:
                        rest_dofs = np.setdiff1d(np.unique(mesh.edges),nodesDBC)
                    elif ndim==3:
                        rest_dofs = np.setdiff1d(np.unique(mesh.faces),nodesDBC)

                    rest_out = np.repeat(rest_dofs,nvar)*nvar + np.tile(np.arange(nvar),rest_dofs.shape[0])
                    rest_app = np.zeros(rest_dofs.shape[0]*nvar)

                    self.columns_out = np.concatenate((self.columns_out,rest_out)).astype(np.int64)
                    self.applied_dirichlet = np.concatenate((self.applied_dirichlet,rest_app))

                print('Finished identifying Dirichlet boundary conditions from CAD geometry.',
                    ' Time taken', time()-tCAD, 'seconds')

            else:

                end = -3
                self.applied_dirichlet = np.loadtxt(mesh.filename.split(".")[0][:end]+"_dirichlet.dat",  dtype=np.float64)
                self.columns_out = np.loadtxt(mesh.filename.split(".")[0][:end]+"_columns_out.dat")

                print('Finished identifying Dirichlet boundary conditions from CAD geometry.',
                    ' Time taken', time()-tCAD, 'seconds')

        #----------------------------------------------------------------------------------------------------#
        #------------------------------------- NON-NURBS BASED SOLUTION -------------------------------------#
        #----------------------------------------------------------------------------------------------------#

        elif self.boundary_type == 'straight' or self.boundary_type == 'mixed':
            # IF DIRICHLET BOUNDARY CONDITIONS ARE APPLIED DIRECTLY AT NODES
            if self.dirichlet_flags is None:
                raise RuntimeError("Dirichlet boundary conditions are not set for the analysis")

            if self.dirichlet_data_applied_at == 'node':
                if self.analysis_type == "dynamic":
                    # FOR DYNAMIC ANALYSIS IT IS ASSUMED THAT
                    # self.columns_in and self.columns_out DO NOT CHANGE
                    # DURING THE ANALYSIS
                    if self.dirichlet_flags.ndim == 3:
                        flat_dirich = self.dirichlet_flags[:,:,0].ravel()
                        self.columns_out = np.arange(self.dirichlet_flags[:,:,0].size)[~np.isnan(flat_dirich)]
                        self.applied_dirichlet = np.zeros((self.columns_out.shape[0],self.dirichlet_flags.shape[2]))

                        for step in range(self.dirichlet_flags.shape[2]):
                            flat_dirich = self.dirichlet_flags[:,:,step].ravel()
                            self.applied_dirichlet[:,step] = flat_dirich[~np.isnan(flat_dirich)]

                    elif self.dirichlet_flags.ndim == 2:
                        flat_dirich = self.dirichlet_flags.ravel()
                        self.columns_out = np.arange(self.dirichlet_flags.size)[~np.isnan(flat_dirich)]
                        self.applied_dirichlet = flat_dirich[~np.isnan(flat_dirich)]
                    else:
                        raise ValueError("Incorrect Dirichlet flags for dynamic analysis")

                else:
                    flat_dirich = self.dirichlet_flags.ravel()
                    self.columns_out = np.arange(self.dirichlet_flags.size)[~np.isnan(flat_dirich)]
                    self.applied_dirichlet = flat_dirich[~np.isnan(flat_dirich)]

        # GENERAL PROCEDURE - GET REDUCED MATRICES FOR FINAL SOLUTION
        self.columns_out = self.columns_out.astype(np.int64)
        self.columns_in = np.delete(np.arange(0,nvar*mesh.points.shape[0]),self.columns_out)

        if self.columns_in.shape[0] == 0:
            warn("No Dirichlet boundary conditions have been applied. The system is unconstrained")
        if self.columns_out.shape[0] == 0:
            warn("Dirichlet boundary conditions have been applied on the entire mesh")

        if self.save_dirichlet_data:
            from scipy.io import savemat
            diri_dict = {'columns_in':self.columns_in,
                'columns_out':self.columns_out,
                'applied_dirichlet':self.applied_dirichlet}
            savemat(self.filename,diri_dict, do_compression=True)

    def ComputeNeumannForces(self, mesh, materials, function_spaces, compute_traction_forces=True, compute_body_forces=False):
        """Compute/assemble traction and body forces"""

        if self.neumann_flags is None:
            return np.zeros((mesh.points.shape[0]*materials[0].nvar,1),dtype=np.float64)

        nvar = materials[0].nvar
        ndim = mesh.InferSpatialDimension()

        if self.neumann_flags.shape[0] == mesh.points.shape[0]:
            self.neumann_data_applied_at = "node"
        else:
            if ndim==3:
                if self.neumann_flags.shape[0] == mesh.faces.shape[0]:
                    self.neumann_data_applied_at = "face"
            elif ndim==2:
                if self.neumann_flags.shape[0] == mesh.edges.shape[0]:
                    self.neumann_data_applied_at = "face"

        if self.neumann_data_applied_at == 'face':
            from Kuru.FiniteElements.Assembly import AssembleForces
            if not isinstance(function_spaces,tuple):
                raise ValueError("Boundary functional spaces not available for computing Neumman and body forces")
            else:
                # CHECK IF A FUNCTION SPACE FOR BOUNDARY EXISTS - SAFEGAURDS AGAINST FORMULATIONS THAT DO NO PROVIDE ONE
                has_boundary_spaces = False
                for fs in function_spaces:
                    if ndim == 3 and fs.ndim == 2:
                        has_boundary_spaces = True
                        break
                    elif ndim == 2 and fs.ndim == 1:
                        has_boundary_spaces = True
                        break
                if not has_boundary_spaces:
                    from Kuru import QuadratureRule, FunctionSpace
                    # COMPUTE BOUNDARY FUNCTIONAL SPACES
                    p = mesh.InferPolynomialDegree()
                    bquadrature = QuadratureRule(optimal=3, norder=2*p+1,
                        mesh_type=mesh.boundary_element_type, is_flattened=False)
                    bfunction_space = FunctionSpace(mesh.CreateDummyLowerDimensionalMesh(),
                        bquadrature, p=p, equally_spaced=mesh.IsEquallySpaced, use_optimal_quadrature=False)
                    function_spaces = (function_spaces[0],bfunction_space)
                    # raise ValueError("Boundary functional spaces not available for computing Neumman and body forces")

            t_tassembly = time()
            if self.analysis_type == "static":
                F = AssembleForces(self, mesh, materials, function_spaces,
                    compute_traction_forces=compute_traction_forces, compute_body_forces=compute_body_forces)
            elif self.analysis_type == "dynamic":
                if self.neumann_flags.ndim==2:
                    # THE POSITION OF NEUMANN DATA APPLIED AT FACES CAN CHANGE DYNAMICALLY
                    tmp_flags = np.copy(self.neumann_flags)
                    tmp_data = np.copy(self.applied_neumann)
                    F = np.zeros((mesh.points.shape[0]*nvar,self.neumann_flags.shape[1]))
                    for step in range(self.neumann_flags.shape[1]):
                        self.neumann_flags = tmp_flags[:,step]
                        self.applied_neumann = tmp_data[:,:,step]
                        F[:,step] = AssembleForces(self, mesh, materials, function_spaces,
                            compute_traction_forces=compute_traction_forces, compute_body_forces=compute_body_forces).flatten()

                    self.neumann_flags = tmp_flags
                    self.applied_neumann = tmp_data
                else:
                    # THE POSITION OF NEUMANN DATA APPLIED AT FACES CAN CHANGE DYNAMICALLY
                    F = AssembleForces(self, mesh, materials, function_spaces,
                            compute_traction_forces=compute_traction_forces, compute_body_forces=compute_body_forces).flatten()

            print("Assembled external traction forces. Time elapsed is {} seconds".format(time()-t_tassembly))

        elif self.neumann_data_applied_at == 'node':
            # A DIRICHLET TYPE METHODOLGY FOR APPLYING NEUMANN BOUNDARY CONDITONS (i.e. AT NODES)
            if self.analysis_type == "dynamic":
                if self.neumann_flags.ndim ==3:
                    # FOR DYNAMIC ANALYSIS IT IS ASSUMED THAT
                    # to_apply DOOES NOT CHANGE DURING THE ANALYSIS
                    flat_neu = self.neumann_flags[:,:,0].ravel()
                    to_apply = np.arange(self.neumann_flags[:,:,0].size)[~np.isnan(flat_neu)]
                    F = np.zeros((mesh.points.shape[0]*nvar,self.neumann_flags.shape[2]))

                    for step in range(self.neumann_flags.shape[2]):
                        flat_neu = self.neumann_flags[:,:,step].ravel()
                        to_apply = np.arange(self.neumann_flags[:,:,step].size)[~np.isnan(flat_neu)]
                        F[to_apply,step] = flat_neu[~np.isnan(flat_neu)]
                else:
                    F = np.zeros((mesh.points.shape[0]*nvar,1))
                    flat_neu = self.neumann_flags.ravel()
                    to_apply = np.arange(self.neumann_flags.size)[~np.isnan(flat_neu)]
                    applied_neumann = flat_neu[~np.isnan(flat_neu)]
                    F[to_apply,0] = applied_neumann
            else:
                F = np.zeros((mesh.points.shape[0]*nvar,1))
                flat_neu = self.neumann_flags.ravel()
                to_apply = np.arange(self.neumann_flags.size)[~np.isnan(flat_neu)]
                applied_neumann = flat_neu[~np.isnan(flat_neu)]
                F[to_apply,0] = applied_neumann

        return F

    def ComputeRobinForces(self, mesh, materials, function_spaces, fem_solver, Eulerx, stiffness, F):
        """Compute/assemble traction and body forces"""

        from Kuru.FiniteElements.Assembly import AssembleRobinForces
        if not self.pressure_flags is None:
            K_pressure, F_pressure = AssembleRobinForces(self, mesh,
                materials[0], function_spaces, fem_solver, Eulerx, 'pressure')
            stiffness -= K_pressure
            F -= F_pressure[:,None]
        if not self.spring_flags is None:
            K_spring, F_spring = AssembleRobinForces(self, mesh,
                materials[0], function_spaces, fem_solver, Eulerx, 'spring')
            stiffness += K_spring
            F += F_spring[:,None]
        if not self.connector_elements is None:
            K_connector, F_connector = AssembleRobinForces(self, mesh,
                materials[0], function_spaces, fem_solver, Eulerx, 'connector')
            stiffness += K_connector
            F += F_connector[:,None]

        return stiffness, F

    def GetReducedMatrices(self, stiffness, F, mass=None, only_residual=False):

        # GET REDUCED FORCE VECTOR
        F_b = F[self.columns_in,0]
        if only_residual:
            return F_b

        # GET REDUCED STIFFNESS MATRIX
        stiffness_b = stiffness[self.columns_in,:][:,self.columns_in]

        # GET REDUCED MASS MATRIX
        mass_b = np.array([])

        return stiffness_b, F_b, mass_b

    def ApplyDirichletGetReducedMatrices(self, stiffness, F, AppliedDirichlet, LoadFactor=1., mass=None, only_residual=False):
        """AppliedDirichlet is a non-member because it can be external incremental Dirichlet,
            which is currently not implemented as member of BoundaryCondition. F also does not
            correspond to Dirichlet forces, as it can be residual in incrementally linearised
            framework.
        """

        # # APPLY DIRICHLET BOUNDARY CONDITIONS
        # for i in range(self.columns_out.shape[0]):
            # F = F - LoadFactor*AppliedDirichlet[i]*stiffness.getcol(self.columns_out[i])

        # MUCH FASTER APPROACH
        # F = F - (stiffness[:,self.columns_out]*AppliedDirichlet*LoadFactor)[:,None]
        nnz_cols = ~np.isclose(AppliedDirichlet,0.0)
        if self.columns_out[nnz_cols].shape[0]==0:
            F[self.columns_in] = F[self.columns_in]
        else:
            F[self.columns_in] = F[self.columns_in] - (stiffness[self.columns_in,:]\
                [:,self.columns_out[nnz_cols]]*AppliedDirichlet[nnz_cols]*LoadFactor)[:,None]

        if only_residual:
            return F

        # GET REDUCED FORCE VECTOR
        F_b = F[self.columns_in,0]

        # GET REDUCED STIFFNESS
        stiffness_b = stiffness[self.columns_in,:][:,self.columns_in]

        # GET REDUCED MASS MATRIX
        if self.analysis_type != 'static':
            mass_b = mass[self.columns_in,:][:,self.columns_in]
            return stiffness_b, F_b, F, mass_b

        return stiffness_b, F_b, F

    def UpdateFixDoFs(self, AppliedDirichletInc, fsize, nvar):
        """Updates the geometry (DoFs) with incremental Dirichlet boundary conditions
            for fixed/constrained degrees of freedom only. Needs to be applied per time steps"""

        # GET TOTAL SOLUTION
        TotalSol = np.zeros((fsize,1))
        TotalSol[self.columns_out,0] = AppliedDirichletInc

        # RE-ORDER SOLUTION COMPONENTS
        dU = TotalSol.reshape(int(TotalSol.shape[0]/nvar),nvar)

        return dU

    def UpdateFreeDoFs(self, sol, fsize, nvar):
        """Updates the geometry with iterative solutions of Newton-Raphson
            for free degrees of freedom only. Needs to be applied per time NR iteration"""

        # GET TOTAL SOLUTION
        TotalSol = np.zeros((fsize,1))
        TotalSol[self.columns_in,0] = sol

        # RE-ORDER SOLUTION COMPONENTS
        dU = TotalSol.reshape(int(TotalSol.shape[0]/nvar),nvar)

        return dU