import sys
import os
import math
import numpy as np
from Parameter import floatParameter, intParameter, boolParameter, OutputParameter, ReadParameter
from Units import *
import Model
from OptionList import ReservoirModel

class WellBores:
    def __init__(self, model:Model):
        """
        The __init__ function is the constructor for a class.  It is called whenever an instance of the class is created.  The __init__ function can take arguments, but self is always the first one. Self refers to the instance of the object that has already been created and it's used to access variables that belong to that object."

        :param self: Reference the class object itself
        :param model: The container class of the application, giving access to everything else, including the logger

        :return: Nothing, and is used to initialize the class
        :doc-author: Malcolm Ross
        """
        model.logger.info("Init " + str(__class__) + ": " + sys._getframe().f_code.co_name)

        #Set up all the Parameters that will be predefined by this class using the different types of parameter classes.  Setting up includes giving it a name, a default value, The Unit Type (length, volume, temperature, etc) and Unit Name of that value, sets it as required (or not), sets allowable range, the error message if that range is exceeded, the ToolTip Text, and the name of teh class that created it.
        #This includes setting up temporary variables that will be available to all the class but noy read in by user, or used for Output
        #This also includes all Parameters that are calculated and then published using the Printouts function.
        #If you choose to sublass this master class, you can do so before or after you create your own parameters.  If you do, you can also choose to call this method from you class, which will effectively add and set all these parameters to your class.

        #These disctionarie contains a list of all the parameters set in this object, stored as "Parameter" and OutputParameter Objects.  This will alow us later to access them in a user interface and get that list, along with unit type, preferred units, etc.
        self.ParameterDict = {}
        self.OutputParameterDict = {}

        self.nprod = self.ParameterDict[self.nprod.Name] = intParameter("Number of Production Wells", value = 2, DefaultValue=2, AllowableRange=list(range(1,201,1)), UnitType = Units.NONE, Required=True, ErrMessage = "assume default number of production wells (2)", ToolTipText="Number of (identical) production wells")
        self.ninj = self.ParameterDict[self.ninj.Name] = intParameter("Number of Injection Wells", value = 2, DefaultValue=2, AllowableRange=list(range(0,201,1)), UnitType = Units.NONE, Required=True, ErrMessage = "assume default number of injection wells (2)", ToolTipText="Number of (identical) injection wells")
        self.prodwelldiam =  self.ParameterDict[self.prodwelldiam.Name] = floatParameter("Production Well Diameter", value = 8.0, DefaultValue=8.0, Min=1, Max=30, UnitType = Units.LENGTH, PreferredUnits = LengthUnit.INCHES, CurrentUnits = LengthUnit.INCHES, Required=True, ErrMessage = "assume default production well diameter (8 inch)", ToolTipText="Inner diameter of production wellbore (assumed constant along the wellbore) to calculate frictional pressure drop and wellbore heat transmission with Rameys model")
        self.injwelldiam = self.ParameterDict[self.injwelldiam.Name] = floatParameter("Injection Well Diameter", value = 8.0, DefaultValue=8.0, Min=1, Max=30, UnitType = Units.LENGTH, PreferredUnits = LengthUnit.INCHES, CurrentUnits = LengthUnit.INCHES, Required=True, ErrMessage = "assume default injection well diameter (8 inch)", ToolTipText="Inner diameter of production wellbore (assumed constant along the wellbore) to calculate frictional pressure drop and wellbore heat transmission with Rameys model")
        self.rameyoptionprod = self.ParameterDict[self.rameyoptionprod.Name] = boolParameter("Ramey Production Wellbore Model", value = True, DefaultValue=True, UnitType = Units.NONE, Required=True, ErrMessage = "assume default production wellbore model (Ramey model active)", ToolTipText="Select whether to use Rameys model to estimate the geofluid temperature drop in the production wells")
        self.tempdropprod = self.ParameterDict[self.tempdropprod.Name] = floatParameter("Production Wellbore Temperature Drop", value = 5.0, DefaultValue=5.0, Min=-5, Max=50, UnitType = Units.TEMPERATURE, PreferredUnits = TemperatureUnit.CELCIUS, CurrentUnits = TemperatureUnit.CELCIUS, ErrMessage = "assume default production wellbore temperature drop (5 deg.C)", ToolTipText="Specify constant production well geofluid temperature drop in case Rameys model is disabled.")
        self.tempgaininj = self.ParameterDict[self.tempgaininj.Name] = floatParameter("Injection Wellbore Temperature Gain", value = 0.0, DefaultValue=0.0, Min=-5, Max=50, UnitType = Units.TEMPERATURE, PreferredUnits = TemperatureUnit.CELCIUS, CurrentUnits = TemperatureUnit.CELCIUS, ErrMessage = "assume default injection wellbore temperature gain (0 deg.C)", ToolTipText="Specify constant injection well geofluid temperature gain.")
        self.prodwellflowrate = self.ParameterDict[self.prodwellflowrate.Name] = floatParameter("Production Flow Rate per Well", value = 50.0, DefaultValue=50.0, Min = 1, Max=500, UnitType = Units.FLOWRATE, PreferredUnits = FlowRateUnit.KGPERSEC, CurrentUnits = FlowRateUnit.KGPERSEC, ErrMessage = "assume default flow rate per production well (50 kg/s)", ToolTipText="Geofluid flow rate per production well.")
        self.impedance = self.ParameterDict[self.impedance.Name] = floatParameter("Reservoir Impedance", value = 1000.0, DefaultValue=1000.0, Min = 1E-4, Max = 1E4, UnitType = Units.IMPEDANCE, PreferredUnits = ImpedanceUnit.GPASPERM3, CurrentUnits = ImpedanceUnit.GPASPERM3, ErrMessage = "assume default reservoir impedance (0.1 GPa*s/m^3)", ToolTipText="Reservoir resistance to flow per well-pair. For EGS-type reservoirs when the injection well is in hydraulic communication with the production well, this parameter specifies the overall pressure drop in the reservoir between injection well and production well (see docs)")
        self.wellsep = self.ParameterDict[self.wellsep.Name] = floatParameter("Well Separation", value = 1000.0, DefaultValue=1000.0, Min=10, Max=10000, UnitType = Units.LENGTH, PreferredUnits = LengthUnit.METERS, CurrentUnits = LengthUnit.INCHES, ErrMessage = "assume default well seperation (1000 m)", ToolTipText="Well separation for built-in TOUGH2 doublet reservoir model")
        self.Tinj = self.ParameterDict[self.Tinj.Name] = floatParameter("Injection Temperature", value = 70.0, DefaultValue=70.0, Min= 0, Max=200, UnitType = Units.TEMPERATURE, PreferredUnits = TemperatureUnit.CELCIUS, CurrentUnits = TemperatureUnit.CELCIUS, Required=True, ErrMessage = "assume default injection temperature (70 deg.C)", ToolTipText="Constant geofluid injection temperature at injection wellhead.")
        self.Phydrostatic = self.ParameterDict[self.Phydrostatic.Name] = floatParameter("Reservoir Hydrostatic Pressure", value = 1E2, DefaultValue=1E2, Min=1E2, Max=1E5, UnitType = Units.PRESSURE, PreferredUnits = PressureUnit.KPASCAL, CurrentUnits = PressureUnit.KPASCAL, ErrMessage = "calculate reservoir hydrostatic pressure using built-in correlation", ToolTipText="Reservoir hydrostatic far-field pressure.  Default value is calculated with built-in modified Xie-Bloomfield-Shook equation (DOE, 2016).")
        self.ppwellhead = self.ParameterDict[self.ppwellhead.Name] = floatParameter("Production Wellhead Pressure", value = 101.3200 + 344.7, DefaultValue=101.3200 + 344.7, Min=0.0, Max=1E4, UnitType = Units.PRESSURE, PreferredUnits = PressureUnit.KPASCAL, CurrentUnits = PressureUnit.KPASCAL, ErrMessage = "using Water vapor pressure (101.3200; at 100 C) + 344.7 kPa (50 psi).", ToolTipText="Constant production wellhead pressure; Required if specifying productivity index")
        self.II = self.ParameterDict[self.II.Name] = floatParameter("Injectivity Index", value = 10.0, DefaultValue=10.0, Min=1E-2, Max=1E4, UnitType = Units.INJECTIVITY_INDEX, PreferredUnits = InjectivityIndexUnit.KGPERSECPERBAR, CurrentUnits = InjectivityIndexUnit.KGPERSECPERBAR, ErrMessage = "assume default injectivity index (10 kg/s/bar)", ToolTipText="Injectivity index defined as ratio of injection well flow rate over injection well outflow pressure drop (flowing bottom hole pressure - hydrostatic reservoir pressure).")
        self.PI = self.ParameterDict[self.PI.Name] = floatParameter("Productivity Index", value = 10.0, DefaultValue=10.0, Min=1E-2, Max=1E4, UnitType = Units.PRODUCTIVITY_INDEX, PreferredUnits = ProductivityIndexUnit.KGPERSECPERBAR, CurrentUnits = ProductivityIndexUnit.KGPERSECPERBAR, ErrMessage = "assume default productivity index (10 kg/s/bar)", ToolTipText="Productivity index defined as ratio of production well flow rate over production well inflow pressure drop (see docs)")
        self.maxdrawdown = self.ParameterDict[self.maxdrawdown.Name] = floatParameter("Maximum Drawdown", value = 1.0, DefaultValue=1.0, Min=0.0, Max=1.000001, UnitType = Units.PERCENT, PreferredUnits = PercentUnit.TENTH, CurrentUnits = PercentUnit.TENTH, ErrMessage = "assume default maximum drawdown (1)", ToolTipText="Maximum allowable thermal drawdown before redrilling of all wells into new reservoir (most applicable to EGS-type reservoirs with heat farming strategies). E.g. a value of 0.2 means that all wells are redrilled after the production temperature (at the wellhead) has dropped by 20% of its initial temperature")
#        self.HasHorizontalSection = self.ParameterDict[self.HasHorizontalSection.Name] = boolParameter("Has Nonvertical Section", value = False, DefaultValue=False, UnitType = Units.NONE, Required=False, ErrMessage = "assume default has no horizontal section", ToolTipText="Set to true if the wells have a horizontal section (as in closed loops)")
        self.IsAGS = self.ParameterDict[self.IsAGS.Name] = boolParameter("Is AGS", value = False, DefaultValue=False, UnitType = Units.NONE, Required=False, ErrMessage = "assume default is not AGS", ToolTipText="Set to true if the model is for an Advanced Geothermal System (AGS)")

#MIR Does Tinj, ninj, nprod, prodwellflowrate change values?

        #local variable initiation
        self.Pinjwellhead = 0.0
        self.usebuiltinhydrostaticpressurecorrelation = True
        self.usebuiltinppwellheadcorrelation = True
        self.Pminimum = 0.0
        sclass = str(__class__).replace("<class \'", "")
        self.MyClass = sclass.replace("\'>","")
        self.MyPath = os.path.abspath(__file__)

        #Results - used by other objects or printed in output downstream
        self.Phydrostaticcalc = self.OutputParameterDict[self.Phydrostaticcalc.Name] = floatParameter("Calculated Reservoir Hydrostatic Pressure", value = self.Phydrostatic.value, UnitType = Units.PRESSURE, PreferredUnits = PressureUnit.KPASCAL, CurrentUnits = PressureUnit.KPASCAL)
        self.redrill = self.OutputParameterDict[self.redrill.Name] = OutputParameter(Name = "redrill", value=0, UnitType = Units.NONE)
        self.PumpingPowerProd = self.OutputParameterDict[self.PumpingPowerProd.Name] = OutputParameter(Name = "PumpingPowerProd", value=[0.0], UnitType = Units.POWER, PreferredUnits = PowerUnit.MW, CurrentUnits = PowerUnit.MW)
        self.PumpingPowerInj = self.OutputParameterDict[self.PumpingPowerInj.Name] = OutputParameter(Name = "PumpingPowerInj", value=[0.0], UnitType = Units.POWER, PreferredUnits = PowerUnit.MW, CurrentUnits = PowerUnit.MW)
        self.pumpdepth = self.OutputParameterDict[self.pumpdepth.Name] = OutputParameter(Name = "pumpdepth", value=[0.0], UnitType = Units.LENGTH, PreferredUnits = LengthUnit.METERS, CurrentUnits = LengthUnit.METERS)
        self.impedancemodelallowed = self.OutputParameterDict[self.impedancemodelallowed.Name] = OutputParameter(Name = "impedancemodelallowed", value=True, UnitType = Units.NONE)
        self.productionwellpumping = self.OutputParameterDict[self.productionwellpumping.Name] = OutputParameter(Name = "productionwellpumping", value=True, UnitType = Units.NONE)
        self.impedancemodelused = self.OutputParameterDict[self.impedancemodelused.Name] = OutputParameter(Name = "impedancemodelused", value=False, UnitType = Units.NONE)
        self.ProdTempDrop = self.OutputParameterDict[self.ProdTempDrop.Name] = OutputParameter(Name = "Production Well Temperature Drop", value=[0.0], UnitType = Units.TEMPERATURE, PreferredUnits = TemperatureUnit.CELCIUS, CurrentUnits = TemperatureUnit.CELCIUS)
        self.DPOverall = self.OutputParameterDict[self.DPOverall.Name] = OutputParameter(Name = "Total Pressure Drop", value=[0.0], UnitType = Units.PRESSURE, PreferredUnits = PressureUnit.KPASCAL, CurrentUnits = PressureUnit.KPASCAL)
        self.DPInjWell = self.OutputParameterDict[self.DPInjWell.Name] = OutputParameter(Name = "Injection Well Pressure Drop", value=[0.0], UnitType = Units.PRESSURE, PreferredUnits = PressureUnit.KPASCAL, CurrentUnits = PressureUnit.KPASCAL)
        self.DPReserv = self.OutputParameterDict[self.DPReserv.Name] = OutputParameter(Name = "Reservoir Pressure Drop", value=[0.0], UnitType = Units.PRESSURE, PreferredUnits = PressureUnit.KPASCAL, CurrentUnits = PressureUnit.KPASCAL)
        self.DPProdWell = self.OutputParameterDict[self.DPProdWell.Name] = OutputParameter(Name = "Production Well Pump Pressure Drop", value=[0.0], UnitType = Units.PRESSURE, PreferredUnits = PressureUnit.KPASCAL, CurrentUnits = PressureUnit.KPASCAL)
        self.DPBouyancy = self.OutputParameterDict[self.DPBouyancy.Name] = OutputParameter(Name = "Bouyancy Pressure Drop", value=[0.0], UnitType = Units.PRESSURE, PreferredUnits = PressureUnit.KPASCAL, CurrentUnits = PressureUnit.KPASCAL)
        self.ProducedTemperature = self.OutputParameterDict[self.ProducedTemperature.Name] = OutputParameter(Name = "Produced Temperature", value=[0.0], UnitType = Units.TEMPERATURE, PreferredUnits = TemperatureUnit.CELCIUS, CurrentUnits = TemperatureUnit.CELCIUS)
        self.PumpingPower = self.OutputParameterDict[self.PumpingPower.Name] = OutputParameter(Name = "Pumping Power", value=[0.0], UnitType = Units.POWER, PreferredUnits = PowerUnit.MW, CurrentUnits = PowerUnit.MW)
        self.Pprodwellhead = self.OutputParameterDict[self.Pprodwellhead.Name] = OutputParameter(Name = "Production wellhead pressure", value=-999.0, UnitType = Units.PRESSURE, PreferredUnits = PressureUnit.KPASCAL, CurrentUnits = PressureUnit.KPASCAL)

        model.logger.info("Complete "+ str(__class__) + ": " + sys._getframe().f_code.co_name)

    def __str__(self):
        return "WellBores"

    def read_parameters(self, model:Model) -> None:
        """
        The read_parameters function reads in the parameters from a dictionary and stores them in the aparmeters.  It also handles special cases that need to be handled after a value has been read in and checked.  If you choose to sublass this master class, you can also choose to override this method (or not), and if you do

        :param self: Access variables that belong to a class
        :param model: The container class of the application, giving access to everything else, including the logger

        :return: None
        :doc-author: Malcolm Ross
        """
        model.logger.info("Init " + str(__class__) + ": " + sys._getframe().f_code.co_name)

        #Deal with all the parameter values that the user has provided.  They should really only provide values that they want to change from the default values, but they can provide a value that is already set because it is a defaulr value set in __init__.  It will ignore those.
        #This also deals with all the special cases that need to be talen care of after a vlaue has been read in and checked.
        #If you choose to sublass this master class, you can also choose to override this method (or not), and if you do, do it before or after you call you own version of this method.  If you do, you can also choose to call this method from you class, which can effectively modify all these superclass parameters in your class.

        if len(model.InputParameters) > 0:
            #loop thru all the parameters that the user wishes to set, looking for parameters that match this object
            for item in self.ParameterDict.items():
                ParameterToModify = item[1]
                key = ParameterToModify.Name.strip()
                if key in model.InputParameters:
                    ParameterReadIn = model.InputParameters[key]
                    ParameterToModify.CurrentUnits = ParameterToModify.PreferredUnits    #Before we change the paremeter, let's assume that the unit preferences will match - if they don't, the later code will fix this.
                    ReadParameter(ParameterReadIn, ParameterToModify, model)   #this should handle all the non-special cases

                    #handle special cases
                    if ParameterToModify.Name == "Reservoir Impedance":  #impedance: impedance per wellpair (input as GPa*s/m^3 and converted to KPa/kg/s (assuming 1000 for density; density will be corrected for later))
                        self.impedance.value = self.impedance.value * (1E6/1E3)    #shift it by a constant to make the units right, per line 619 of GEOPHIRES 2
                        self.impedancemodelused.value = True
                        if self.impedance.Provided == False: self.impedancemodelused.value = False
                    elif ParameterToModify.Name == "Reservoir Hydrostatic Pressure":
                        if ParameterToModify.value == -1: self.usebuiltinhydrostaticpressurecorrelation = True
                        else: self.usebuiltinhydrostaticpressurecorrelation = False
                    elif ParameterToModify.Name == "Production Wellhead Pressure":
                        if ParameterToModify.value == -1.0: self.usebuiltinppwellheadcorrelation = True
                        else: self.usebuiltinppwellheadcorrelation = False
        else:
            model.logger.info("No parameters read becuase no content provided")
        model.logger.info("read parameters complete "+ str(__class__) + ": " + sys._getframe().f_code.co_name)

    #user-defined functions
    def vaporpressurewater(self, Twater) -> float:
        if Twater < 100:
            A = 8.07131
            B = 1730.63
            C = 233.426
        else:
            A = 8.14019
            B = 1810.94
            C = 244.485
        vaporpressurewater = 133.322*(10**(A-B/(C+Twater)))/1000 #water vapor pressure in kPa using Antione Equation
        return vaporpressurewater

    def RameyCalc(self, krock:float, rhorock:float, cprock:float, welldiam:float, tv, utilfactor:float, flowrate:float, cpwater:float, Trock:float,
                  Tresoutput:float, averagegradient:float, depth:float):
        #this code is only valid so far for 1 gradient and deviation = 0
        #For multiple gradients, use Ramey's model for every layer
        #assume outside diameter of casing is 10% larger than inside diameter of production pipe (=prodwelldiam)
        #assume borehole thermal resistance negligible to rock thermal resistance
        alen = len(tv)
        alpharock = krock/(rhorock*cprock)
        framey = np.zeros(alen)
        framey[1:] = -np.log(1.1*(welldiam/2.0)/np.sqrt(4.*alpharock*tv[1:]*365.0*24.0*3600.0*utilfactor))-0.29
        framey[0] = framey[1]    #fource the first value to be the same as the second to get away from near surface effects
        rameyA = flowrate*cpwater*framey/2/math.pi/krock
        TempDrop = -((Trock - Tresoutput) - averagegradient*(depth - rameyA) + (Tresoutput - averagegradient*rameyA - Trock)*np.exp(-depth/rameyA))

        return TempDrop

    def WellPressureDrop(self, model:Model, Taverage:float, wellflowrate:float, welldiam:float, impedancemodelused:bool, depth:float):
        #start by calculating wellbore fluid conditions [kPa], noting that most of temperature drop happens in upper section (because surrounding rock temperature is lowest in upper section)
        rhowater = model.reserv.densitywater(Taverage)  #replace with correlation based on Tprodaverage
        muwater = model.reserv.viscositywater(Taverage) #replace with correlation based on Tprodaverage
        v = wellflowrate/rhowater/(math.pi/4.*welldiam**2)
        Rewater = 4.0*wellflowrate/(muwater*math.pi*welldiam) #laminar or turbulent flow?
        Rewateraverage = np.average(Rewater)
        if Rewateraverage < 2300.0: f3 = 64./Rewater
        else:
            relroughness = 1E-4/welldiam
            f3 = 1./np.power(-2*np.log10(relroughness/3.7+5.74/np.power(Rewater,0.9)),2.)
            f3 = 1./np.power((-2*np.log10(relroughness/3.7+2.51/Rewater/np.sqrt(f3))),2.)
            f3 = 1./np.power((-2*np.log10(relroughness/3.7+2.51/Rewater/np.sqrt(f3))),2.)
            f3 = 1./np.power((-2*np.log10(relroughness/3.7+2.51/Rewater/np.sqrt(f3))),2.)
            f3 = 1./np.power((-2*np.log10(relroughness/3.7+2.51/Rewater/np.sqrt(f3))),2.)
            f3 = 1./np.power((-2*np.log10(relroughness/3.7+2.51/Rewater/np.sqrt(f3))),2.) #6 iterations to converge

        if impedancemodelused: DPWell = f3*(rhowater*v**2/2.)*(depth/welldiam)/1E3     ##assumed everything stays liquid throughout; /1E3 to convert from Pa to kPa
        else: DPWell = []    #it will be calculated elsewhere

        return DPWell, f3, v, rhowater

    def InjectionWellPressureDrop(self, model:Model, Taverage:float, wellflowrate:float, welldiam:float, impedancemodelused:bool, depth:float, nprod:int, ninj:int, waterloss:float):
        #start by calculating wellbore fluid conditions [kPa], noting that most of temperature drop happens in upper section (because surrounding rock temperature is lowest in upper section)
        rhowater = model.reserv.densitywater(Taverage)*np.linspace(1,1,len(self.ProducedTemperature.value))
        muwater = model.reserv.viscositywater(Taverage)*np.linspace(1,1,len(self.ProducedTemperature.value))  #replace with correlation based on Tinjaverage
        v = nprod/ninj*wellflowrate*(1.0 + waterloss)/rhowater/(math.pi/4.*welldiam**2)
        Rewater = 4.*nprod/ninj*wellflowrate*(1.0 + waterloss)/(muwater*math.pi*welldiam) #laminar or turbulent flow?
        Rewateraverage = np.average(Rewater)
        if Rewateraverage < 2300. : #laminar flow
            f1 = 64./Rewater
        else: #turbulent flow
            relroughness = 1E-4/welldiam
            f1 = 1./np.power(-2*np.log10(relroughness/3.7+5.74/np.power(Rewater,0.9)),2.)
            f1 = 1./np.power((-2*np.log10(relroughness/3.7+2.51/Rewater/np.sqrt(f1))),2.)
            f1 = 1./np.power((-2*np.log10(relroughness/3.7+2.51/Rewater/np.sqrt(f1))),2.)
            f1 = 1./np.power((-2*np.log10(relroughness/3.7+2.51/Rewater/np.sqrt(f1))),2.)
            f1 = 1./np.power((-2*np.log10(relroughness/3.7+2.51/Rewater/np.sqrt(f1))),2.)
            f1 = 1./np.power((-2*np.log10(relroughness/3.7+2.51/Rewater/np.sqrt(f1))),2.)  #6 iterations to converge

        if self.impedancemodelused.value: DPWell = f1*(rhowater*v**2/2)*(depth/welldiam)/1E3    #assumed everything stays liquid throughout, calculate well pressure drop [kPa]; /1E3 to convert from Pa to kPa
        else: DPWell = []    #it will be calculated elsewhere

        return DPWell, f1, v, rhowater

    def ProdPressureDropsAndPumpingPowerUsingImpedenceModel(self, f3:float, vprod:float, rhowaterinj:float, rhowaterprod:float,
                                                        rhowaterreservoir:float, depth:float, wellflowrate:float, prodwelldiam:float,
                                                        impedance:float, nprod:int, waterloss:float, pumpeff:float):
        #Calculate Pressure Drops and Pumping Power needed for the production well using the Impedance Model
        #production well pressure drops [kPa]
        DPProdWell = f3*(rhowaterprod*vprod**2/2.)*(depth/prodwelldiam)/1E3     #/1E3 to convert from Pa to kPa

        #reservoir pressure drop [kPa]
        DPReserv = impedance*nprod*wellflowrate*1000./rhowaterreservoir

        #buoyancy pressure drop [kPa]
        DPBouyancy = (rhowaterprod-rhowaterinj)*depth*9.81/1E3 # /1E3 to convert from Pa to kPa

        #overall pressure drop
        DPOverall = DPReserv + DPProdWell + DPBouyancy

        #calculate pumping power [MWe] (approximate)
        PumpingPower = DPOverall*nprod*wellflowrate*(1+waterloss)/rhowaterinj/pumpeff/1E3

        #in GEOPHIRES v1.2, negative pumping power values become zero (b/c we are not generating electricity)
        PumpingPower = [0. if x<0. else x for x in PumpingPower]

        return DPOverall, PumpingPower, DPProdWell, DPReserv, DPBouyancy

    def InjPressureDropsAndPumpingPowerUsingImpedenceModel(self, f1:float, vinj:float, rhowaterinj:float, depth:float, wellflowrate:float,
                                                          injwelldiam:float, ninj:int, waterloss:float, pumpeff:float, DPOverall):
        #Calculate Pressure Drops and Pumping Power needed for the injection well using the Impedance Model
        #injection well pressure drops [kPa]
        DPInjWell = f1*(rhowaterinj*vinj**2/2)*(depth/injwelldiam)/1E3      #/1E3 to convert from Pa to kPa

        #overall pressure drop
        newDPOverall = DPOverall + DPInjWell

        #calculate pumping power [MWe] (approximate)
        PumpingPower = newDPOverall*ninj*wellflowrate*(1+waterloss)/rhowaterinj/pumpeff/1E3

        #in GEOPHIRES v1.2, negative pumping power values become zero (b/c we are not generating electricity)
        PumpingPower = [0. if x<0. else x for x in PumpingPower]

        return newDPOverall, PumpingPower, DPInjWell

    def ProdPressureDropAndPumpingPowerUsingIndexes(self, model:Model, usebuiltinhydrostaticpressurecorrelation:bool, productionwellpumping:bool, usebuiltinppwellheadcorrelation:bool,
                                                    Trock:float, Tsurf:float, depth:float, gradient:float, ppwellhead:float, PI: float, wellflowrate:float, f3:float, vprod:float,
                                                    prodwelldiam:float, nprod:int, pumpeff:float,  rhowaterprod:float):
        #Calculate PressureDrops and Pumping Power needed for the production well using indexes
        PumpingPower = PumpingPowerProd = DPProdWell = Pprodwellhead = ([0.0] * len(vprod))  #initialize value in case it doesn't get set.
        #reservoir hydrostatic pressure [kPa]
        if usebuiltinhydrostaticpressurecorrelation:
            CP = 4.64E-7
            CT = 9E-4/(30.796*Trock**(-0.552))
            Phydrostaticcalc = 0+1./CP*(math.exp(model.reserv.densitywater(Tsurf)*9.81*CP/1000*(depth-CT/2*gradient*depth**2))-1)

        if productionwellpumping:
            Pexcess = 344.7 #[kPa] = 50 psi. Excess pressure covers non-condensable gas pressure and net positive suction head for the pump
            Pminimum = self.vaporpressurewater(Trock) + Pexcess #[kPa] is minimum production pump inlet pressure and minimum wellhead pressure
            if usebuiltinppwellheadcorrelation: Pprodwellhead = Pminimum   #production wellhead pressure [kPa]
            else:
                Pprodwellhead = ppwellhead
                if Pprodwellhead < Pminimum:
                    Pprodwellhead = Pminimum
                    print("Warning: provided production wellhead pressure under minimum pressure. GEOPHIRES will assume minimum wellhead pressure")
                    model.logger.warning("Provided production wellhead pressure under minimum pressure. GEOPHIRES will assume minimum wellhead pressure")

            PIkPa = PI/100.0 #convert PI from kg/s/bar to kg/s/kPa

            #calculate pumping depth
            pumpdepth = depth + (Pminimum - Phydrostaticcalc + wellflowrate/PIkPa)/(f3*(rhowaterprod*vprod**2/2.)*(1/prodwelldiam)/1E3 + rhowaterprod*9.81/1E3)
            pumpdepthfinal = np.max(pumpdepth)
            if pumpdepthfinal < 0.0:
                pumpdepthfinal = 0.0
                print("Warning: GEOPHIRES calculates negative production well pumping depth. No production well pumps will be assumed")
                model.logger.warning("GEOPHIRES calculates negative production well pumping depth. No production well pumps will be assumed")
            elif pumpdepthfinal > 600.0:
                print("Warning: GEOPHIRES calculates production pump depth to be deeper than 600 m. Verify reservoir pressure, production well flow rate and production well dimensions")
                model.logger.warning("GEOPHIRES calculates production pump depth to be deeper than 600 m. Verify reservoir pressure, production well flow rate and production well dimensions")

            #calculate production well pumping pressure [kPa]
            DPProdWell = Pprodwellhead - (Phydrostaticcalc - wellflowrate/PIkPa - rhowaterprod*9.81*depth/1E3 - f3*(rhowaterprod*vprod**2/2.)*(depth/prodwelldiam)/1E3)
            PumpingPowerProd = DPProdWell*nprod*wellflowrate/rhowaterprod/pumpeff/1E3 #[MWe] total pumping power for production wells
            PumpingPowerProd = np.array([0. if x<0. else x for x in PumpingPowerProd])

        #total pumping power
        if productionwellpumping: PumpingPower = PumpingPowerProd

        #negative pumping power values become zero (b/c we are not generating electricity)
        PumpingPower = [0. if x<0. else x for x in PumpingPower]

        return PumpingPower, PumpingPowerProd, DPProdWell, Pprodwellhead

    def InjPressureDropAndPumpingPowerUsingIndexes(self, model:Model, usebuiltinhydrostaticpressurecorrelation:bool, productionwellpumping:bool, usebuiltinppwellheadcorrelation:bool, usebuiltinoutletplantcorrelation:bool,
                                                   Trock:float, Tsurf:float, depth:float, gradient:float, ppwellhead:float, II:float, wellflowrate:float, f1:float, vinj:float, injwelldiam:float, nprod:int, ninj:int,
                                                   waterloss:float, pumpeff:float,  rhowaterinj:float, Pplantoutlet:float, PumpingPowerProd):
        #Calculate PressureDrops and Pumping Power needed for the injection well using indexes
        PumpingPowerInj = DPInjWell = Pprodwellhead = [0.0]   #initialize value in case it doesn't get set.

        #reservoir hydrostatic pressure [kPa]
        if usebuiltinhydrostaticpressurecorrelation:
            CP = 4.64E-7
            CT = 9E-4/(30.796*Trock**(-0.552))
            Phydrostaticcalc = 0+1./CP*(math.exp(model.reserv.densitywater(Tsurf)*9.81*CP/1000*(depth-CT/2*gradient*depth**2))-1)

        if productionwellpumping:
            Pexcess = 344.7 #[kPa] = 50 psi. Excess pressure covers non-condensable gas pressure and net positive suction head for the pump
            Pminimum = self.vaporpressurewater(Trock) + Pexcess #[kPa] is minimum production pump inlet pressure and minimum wellhead pressure
            if usebuiltinppwellheadcorrelation: Pprodwellhead = Pminimum   #production wellhead pressure [kPa]
            else:
                Pprodwellhead = ppwellhead
                if Pprodwellhead < Pminimum:
                    Pprodwellhead = Pminimum
                    print("Warning: provided production wellhead pressure under minimum pressure. GEOPHIRES will assume minimum wellhead pressure")
                    model.logger.warning("Provided production wellhead pressure under minimum pressure. GEOPHIRES will assume minimum wellhead pressure")

        IIkPa = II/100.0 #convert II from kg/s/bar to kg/s/kPa

        #necessary injection wellhead pressure [kPa]
        Pinjwellhead = Phydrostaticcalc + wellflowrate*(1+waterloss)*nprod/ninj/IIkPa - rhowaterinj*9.81*depth/1E3 + f1*(rhowaterinj*vinj**2/2)*(depth/injwelldiam)/1E3

        #plant outlet pressure [kPa]
        if usebuiltinoutletplantcorrelation:
            DPSurfaceplant = 68.95 #[kPa] assumes 10 psi pressure drop in surface equipment
            Pplantoutlet = Pprodwellhead - DPSurfaceplant

        #injection pump pressure [kPa]
        DPInjWell = Pinjwellhead-Pplantoutlet
        PumpingPowerInj = DPInjWell*nprod*wellflowrate*(1+waterloss)/rhowaterinj/pumpeff/1E3 #[MWe] total pumping power for injection wells
        PumpingPowerInj = np.array([0. if x<0. else x for x in PumpingPowerInj])

        #total pumping power
        if productionwellpumping: PumpingPower = PumpingPowerInj + PumpingPowerProd
        else: PumpingPower = PumpingPowerInj

        #negative pumping power values become zero (b/c we are not generating electricity)
        PumpingPower = [0. if x<0. else x for x in PumpingPower]

        return PumpingPower, PumpingPowerInj, DPInjWell, Pplantoutlet, Pprodwellhead

    def Calculate(self, model:Model)-> None:
        """
        The Calculate function is where all the calculations are done.
        This function can be called multiple times, and will only recalculate what has changed each time it is called.

        :param self: Access variables that belongs to the class
        :param model: The container class of the application, giving access to everything else, including the logger
        :return: Nothing, but it does make calculations and set values in the model
        :doc-author: Malcolm Ross
        """
        model.logger.info("Init " + str(__class__) + ": " + sys._getframe().f_code.co_name)

        #This is where all the calcualtions are made using all the values that have been set.
        #If you sublcass this class, you can choose to run these calculations before (or after) your calculations, but that assumes you have set all the values that are required for these calculations
        #If you choose to sublass this master class, you can also choose to override this method (or not), and if you do, do it before or after you call you own version of this method.  If you do, you can also choose to call this method from you class, which can effectively run the calculations of the superclass, making all thr values available to your methods. but you had n=betteer have set all the paremeters!

        #special case: production and injection well diameters are input as inches and call calculations assume meters! Check and change if needed, assuming anything > 2 must be talking about inches
        injwelldiam = self.injwelldiam.value
        if self.injwelldiam.value > 2.0: injwelldiam = self.injwelldiam.value * 0.0254
        prodwelldiam = self.prodwelldiam.value
        if self.prodwelldiam.value > 2.0: prodwelldiam = self.prodwelldiam.value * 0.0254

        #calculate wellbore temperature drop
        self.ProdTempDrop.value = self.tempdropprod.value   #if not Ramey, hard code a user-supplied tempweature drop.
        if self.rameyoptionprod.value:
            if hasattr(model.reserv, "InputDepth"): d = model.reserv.InputDepth.value
            else: d = model.reserv.depth.value
            self.ProdTempDrop.value = self.RameyCalc(model.reserv.krock.value, model.reserv.rhorock.value, model.reserv.cprock.value,
                                  prodwelldiam, model.reserv.timevector.value, model.surfaceplant.utilfactor.value, self.prodwellflowrate.value,
                                  model.reserv.cpwater.value, model.reserv.Trock.value, model.reserv.Tresoutput.value, model.reserv.averagegradient.value, d)

        self.ProducedTemperature.value = model.reserv.Tresoutput.value-self.ProdTempDrop.value

        #redrilling
        if model.reserv.resoption.value in [ReservoirModel.MULTIPLE_PARALLEL_FRACTURES, ReservoirModel.LINEAR_HEAT_SWEEP, ReservoirModel.SINGLE_FRACTURE, ReservoirModel.ANNUAL_PERCENTAGE]: #only applies to the built-in analytical reservoir models
            indexfirstmaxdrawdown = np.argmax(self.ProducedTemperature.value<(1-model.wellbores.maxdrawdown.value)*self.ProducedTemperature.value[0])
            if indexfirstmaxdrawdown > 0:   #redrilling necessary
                self.redrill.value = int(np.floor(len(self.ProducedTemperature.value)/indexfirstmaxdrawdown))
                ProducedTemperatureRepeatead = np.tile(self.ProducedTemperature.value[0:indexfirstmaxdrawdown], self.redrill.value+1)
                self.ProducedTemperature.value = ProducedTemperatureRepeatead[0:len(self.ProducedTemperature.value)]

        #------------------------------------------
        #calculate pressure drops and pumping power
        #------------------------------------------
        self.DPProdWell.value, f3, vprod, self.rhowaterprod = self.WellPressureDrop(model, model.reserv.Tresoutput.value-self.ProdTempDrop.value/4.0, self.prodwellflowrate.value, prodwelldiam, self.impedancemodelused.value, model.reserv.depth.value)
        self.DPInjWell.value, f1, vinj, self.rhowaterinj = self.InjectionWellPressureDrop(model, self.Tinj.value, self.prodwellflowrate.value, injwelldiam, self.impedancemodelused.value, model.reserv.depth.value, self.nprod.value, self.ninj.value, model.reserv.waterloss.value)

        if self.impedancemodelused.value: #assumed everything stays liquid throughout
            rhowaterreservoir = model.reserv.densitywater(0.1*self.Tinj.value+0.9*model.reserv.Tresoutput.value)    #based on TARB in Geophires v1.2
            self.DPOverall.value, self.PumpingPower.value, self.DPProdWell.value, self.DPReserv.value, self.DPBouyancy.value = self.ProdPressureDropsAndPumpingPowerUsingImpedenceModel(f3, vprod,
                                          self.rhowaterinj, self.rhowaterprod, model.reserv.rhowater.value, model.reserv.depth.value, self.prodwellflowrate.value, prodwelldiam, self.impedance.value,
                                          self.nprod.value, model.reserv.waterloss.value, model.surfaceplant.pumpeff.value)
            self.DPOverall.value, self.PumpingPower.value, self.DPInjWell.value = self.InjPressureDropsAndPumpingPowerUsingImpedenceModel(f1, vinj,
                                          self.rhowaterinj, model.reserv.depth.value, self.prodwellflowrate.value, injwelldiam,
                                          self.ninj.value, model.reserv.waterloss.value, model.surfaceplant.pumpeff.value, self.DPOverall.value)

        else: #PI and II are used
            self.PumpingPower.value, self.PumpingPowerProd.value, self.DPProdWell.value, self.Pprodwellhead.value =  self.ProdPressureDropAndPumpingPowerUsingIndexes(model, self.usebuiltinhydrostaticpressurecorrelation, self.productionwellpumping.value, self.usebuiltinppwellheadcorrelation,
                                        model.reserv.Trock.value, model.reserv.Tsurf.value, model.reserv.depth.value, model.reserv.averagegradient.value, self.ppwellhead.value, self.PI.value, self.prodwellflowrate.value, f3, vprod,
                                        prodwelldiam, self.nprod.value, model.surfaceplant.pumpeff.value, self.rhowaterprod)
            self.PumpingPower.value, self.PumpingPowerInj.value, self.DPInjWell.value, model.surfaceplant.Pplantoutlet.value, self.Pprodwellhead.value =  self.InjPressureDropAndPumpingPowerUsingIndexes(model, self.usebuiltinhydrostaticpressurecorrelation, self.productionwellpumping.value, self.usebuiltinppwellheadcorrelation, model.surfaceplant.usebuiltinoutletplantcorrelation.value,
                                        model.reserv.Trock.value, model.reserv.Tsurf.value, model.reserv.depth.value, model.reserv.averagegradient.value, self.ppwellhead.value, self.II.value, self.prodwellflowrate.value, f1, vinj,
                                        injwelldiam, self.nprod.value, self.ninj.value, model.reserv.waterloss.value, model.surfaceplant.pumpeff.value, self.rhowaterinj, model.surfaceplant.Pplantoutlet.value, self.PumpingPowerProd.value)

        model.logger.info("complete "+ str(__class__) + ": " + sys._getframe().f_code.co_name)


import os
import cv2
import optics
import argparse
import tfrecord
import numpy as np
import tensorflow as tf
from util import *

class TensorHolographyModel():
    def __init__(self,
                 hologram_params,
                 training_params,
                 model_params,
                 loss_params,
                 path_params,
                 train_dataset_params,
                 test_dataset_params,
                 validate_dataset_params):

        self.hologram_params = hologram_params
        self.training_params = training_params
        self.model_params = model_params
        self.loss_params = loss_params
        self.path_params = path_params
        self.train_dataset_params = train_dataset_params
        self.test_dataset_params = test_dataset_params
        self.validate_dataset_params = validate_dataset_params
        self.model_vars = None

        config = tf.compat.v1.ConfigProto()
        config.gpu_options.allow_growth = True
        self.sess = tf.compat.v1.Session(config=config)

        # build tfrecord if path_params["gen_record"] is true
        if path_params["gen_record"]:
            generator = tfrecord.TFRecordGeneratorforTH(self.path_params["train_output_path"],
                                            self.path_params["labels"],
                                            self.path_params["train_source_paths"])
            generator.generate_record()
            generator.update_record_paths(self.path_params["test_output_path"],
                                            self.path_params["labels"],
                                            self.path_params["test_source_paths"])
            generator.generate_record()
            generator.update_record_paths(self.path_params["validate_output_path"],
                                            self.path_params["labels"],
                                            self.path_params["validate_source_paths"])
            generator.generate_record()
            print("tfrecord generation done!")

    def _build_model_vars(self):
        fw   = np.full((self.model_params["num_layers"]), self.model_params["filter_width"], dtype=int)
        fnum = np.append(np.full((self.model_params["num_layers"]-1), self.model_params["num_filters_per_layer"], dtype=int),
                         self.model_params["output_dim"]*(self.model_params["interleave_rate"]**2))
        model_vars = {}

        for i in range(self.model_params["num_layers"]):
            # first layer
            if i==0:
                in_dim, out_dim = self.model_params["input_dim"] * (self.model_params["interleave_rate"]**2), fnum[i]
            # last layer
            elif i==self.model_params["num_layers"]-1:
                in_dim, out_dim = fnum[i-1] + self.model_params["input_dim"] * (self.model_params["interleave_rate"]**2),  \
                                  self.model_params["output_dim"] * (self.model_params["interleave_rate"]**2)
            else:
                in_dim, out_dim = fnum[i-1], fnum[i]

            model_vars[i] = {'weights':tf_init_weights([fw[i], fw[i], in_dim, out_dim],
                                                        'xavier',
                                                        xavier_params=(in_dim, out_dim),
                                                        r=self.model_params["weight_var_scale"]),
                            'bias':tf.Variable(tf.random.truncated_normal([out_dim],stddev=self.model_params["bias_stddev"]))
                            }

        return model_vars

    def _build_graph(self, x_in, model_vars, data_format='NCHW'):
        layers = {}
        prev_layers = {}

        # build layers
        print("input data:", x_in.shape)
        if self.model_params["renormalize_input"]:
            x_in = x_in - 0.5

        # interleave the input
        if data_format == 'NCHW':
            x_in = tf_interleave_nonnative(self.model_params["interleave_rate"], x_in)
        elif self.model_params["interleave_rate"] != 1:
            # update this for NHWC
            raise Exception('data_format has to be NCHW for interleave')

        # build graph
        for i in range(self.model_params["num_layers"]):
            if i==0:
                prev_layers[i] = x_in
            elif (i<3) or (i%2==0):
                prev_layers[i] = layers[i-1]
            else:
                prev_layers[i] = layers[i-1] + prev_layers[i-2]
                print('(skip connection: %d, %d)'%(i-1, i-3))

            if i == self.model_params["num_layers"]-1:
                prev_layers[i] = tf.concat([prev_layers[i], x_in], axis = 1 if data_format == 'NCHW' else 3)

            if not i == self.model_params["num_layers"]-1:
                layers[i] = self.model_params["activation_func"](
                                tf.layers.batch_normalization(
                                    tf.nn.bias_add(
                                        tf.nn.conv2d(prev_layers[i],model_vars[i]['weights'],strides=[1,1,1,1], padding='SAME', data_format=data_format),
                                        model_vars[i]['bias'],
                                        data_format=data_format
                                    ),
                                    axis=1 if data_format == 'NCHW' else 3
                                )
                            )
                print("layer %d:" % i, layers[i].shape)
            else:
                # last layer
                field = self.model_params["output_activation_func"](
                            tf.layers.batch_normalization(
                                tf.nn.bias_add(
                                    tf.nn.conv2d(prev_layers[i],model_vars[i]['weights'],strides=[1,1,1,1], padding='SAME', data_format=data_format),
                                    model_vars[i]['bias'],
                                    data_format=data_format
                                ),
                                axis=1 if data_format == 'NCHW' else 3
                            ),
                            name='field'
                        )
                if data_format == 'NCHW':
                    field = tf_deinterleave_nonnative(self.model_params["interleave_rate"], field)
                elif self.model_params["interleave_rate"] != 1:
                    raise Exception('data_format has to be NCHW for interleave')

                # normalize amplitude to [0, sqrt(2)]
                # normalize phase to [0, 1]
                # compute complex field, the phase is renormalized to [-pi, pi]
                if data_format=='NCHW':
                    amp = tf.add(field[:,:self.model_params["output_dim"]//2,:,:]*np.sqrt(0.5), np.sqrt(0.5), name="output_field_amp")
                    phs = tf.add(field[:,self.model_params["output_dim"]//2:,:,:]*0.5, 0.5, name="output_field_phs")
                else:
                    amp = tf.strided_slice(field, [0,0,0,0], [1, field.shape[1], field.shape[2], 3], [1,1,1,1], shrink_axis_mask=1)
                    phs = tf.strided_slice(field, [0,0,0,3], [1, field.shape[1], field.shape[2], 6], [1,1,1,1], shrink_axis_mask=1)
                    amp = tf.add(amp*np.sqrt(0.5), np.sqrt(0.5), name="output_field_amp")
                    phs = tf.add(phs*0.5, 0.5, name="output_field_phs")
                out_field = optics.tf_compl_val(amp, (phs-0.5)*2.0*np.pi, name="out_field")

                print("output tensor:", out_field.shape)

        return out_field, amp, phs

    def _get_dataset_iterators(self):
        # train iterator
        extractor = tfrecord.TFRecordExtractorforTH(self.path_params["train_output_path"],
                                                    self.train_dataset_params,
                                                    self.path_params["labels"])
        train_iterator = extractor.build_dataset()
        extractor.update_record_path(self.path_params["test_output_path"],
                                     self.test_dataset_params,
                                     self.path_params["labels"])
        test_iterator = extractor.build_dataset()
        extractor.update_record_path(self.path_params["validate_output_path"],
                                     self.validate_dataset_params,
                                     self.path_params["labels"])
        validate_iterator = extractor.build_dataset()
        return train_iterator, test_iterator, validate_iterator

    def _preprocess_input(self, example):
        # load RGB-D input
        rgbd = tf.concat([example["img"], example["depth"]], 1)

        # create complex hologram
        holo = optics.tf_compl_val(example["amp"], (example["phs"]-0.5) * 2 * np.pi)
        return rgbd, holo, example["amp"], example["phs"]

    def _setup_train(self):
        # get train and test dataset handle
        train_iterator, test_iterator, validate_iterator = self._get_dataset_iterators()
        train_handle = self.sess.run(train_iterator.string_handle())
        test_handle  = self.sess.run(test_iterator.string_handle())
        validate_handle  = self.sess.run(validate_iterator.string_handle())

        # create feedable handle
        handle   = tf.compat.v1.placeholder(tf.string, shape=[])
        iterator = tf.compat.v1.data.Iterator.from_string_handle(handle, output_types=tf.compat.v1.data.get_output_types(train_iterator))
        example  = iterator.get_next()

        # get data placeholder
        rgbd, holo_in, amp_in, phs_in = self._preprocess_input(example)

        # build model
        self.model_vars = self._build_model_vars()
        holo_out, amp_out, phs_out = self._build_graph(rgbd, self.model_vars)

        return train_handle, test_handle, validate_handle, handle, rgbd, holo_in, amp_in, phs_in, holo_out, amp_out, phs_out

    def validate(self):
        _, _, validate_handle, handle, rgbd, holo_in, amp_in, phs_in, holo_out, amp_out, phs_out = self._setup_train()

        # compute psnr and ssim on the amplitude map
        ssim_amp_loss = tf.reduce_mean(tf.image.ssim(tf.transpose(amp_out, [0,2,3,1]), tf.transpose(amp_in, [0,2,3,1]), 1.0))
        psnr_amp_loss = tf.reduce_mean(tf.image.psnr(tf.transpose(amp_out, [0,2,3,1]), tf.transpose(amp_in, [0,2,3,1]), 1.0))

        # create model saver
        self.saver = tf.compat.v1.train.Saver(max_to_keep=5)

        # initialize variables
        init_op = tf.compat.v1.global_variables_initializer()
        self.sess.run(init_op)

        # reload pre-trained model
        if self.training_params["restore_trained_model"]:
            ckpt = tf.train.get_checkpoint_state(self.path_params["ckpt_parent_path"])
            if ckpt and ckpt.model_checkpoint_path:
                self.saver.restore(self.sess, tf.train.latest_checkpoint(self.path_params["ckpt_parent_path"]))
                print("model restored from %s" % self.path_params["ckpt_parent_path"])
            else:
                print("ERROR: NO RESTORED MODEL...")

        # test on validation dataset
        validate_steps = self.validate_dataset_params["sample_count"] // self.validate_dataset_params["batch"]
        avg_ssim_amp_loss = []
        avg_psnr_amp_loss = []
        for v_step in range(validate_steps):
            print("validate step %d/%d\n" % (v_step, validate_steps))
            ssim_amp_loss_val, psnr_amp_loss_val \
                 = self.sess.run([ssim_amp_loss, psnr_amp_loss], feed_dict={handle: validate_handle})
            avg_ssim_amp_loss.append(ssim_amp_loss_val)
            avg_psnr_amp_loss.append(psnr_amp_loss_val)

        print("validation results: ssim amp loss = %f(%f; %f/%f), pnsr amp loss = %f(%f; %f/%f)\n" % \
            (np.mean(avg_ssim_amp_loss), np.std(avg_ssim_amp_loss), np.amax(avg_ssim_amp_loss), np.amin(avg_ssim_amp_loss),
             np.mean(avg_psnr_amp_loss), np.std(avg_psnr_amp_loss), np.amax(avg_psnr_amp_loss), np.amin(avg_psnr_amp_loss)))

    def evaluate(self, eval_params):
        # define placeholder for input rgbd images
        rgbd = tf.compat.v1.placeholder(tf.float32, [1, self.model_params["input_dim"], eval_params['res_h'], eval_params['res_w']])

        # define wave propagator
        propagator = optics.tf_propagator(
            (eval_params['res_h'], eval_params['res_w']),
            self.hologram_params["pitch"],
            self.hologram_params["wavelengths"],
            method="as",
            double_pad=eval_params['double_pad']
        )

        # build network
        self.model_vars = self._build_model_vars()
        _, amp_out, phs_out = self._build_graph(rgbd, self.model_vars)

        # center phase for each color channel
        phs_out = phs_out - tf.reduce_mean(phs_out, [2,3], keepdims=True) + 0.5

        # compute phase-only hologram
        cpx_out = optics.tf_compl_val(amp_out, (phs_out-0.5) * 2.0 * np.pi)
        if eval_params["use_maimone_dpm"]:
            phs_only, amp_max = optics.tf_dpm_maimone(cpx_out,
                                                      propagator,
                                                      depth_shift=eval_params["depth_shift"],
                                                      adaptive_phs_shift=eval_params['adaptive_phs_shift'],
                                                      batch=1,
                                                      num_channels=3,
                                                      res_h=eval_params['res_h'],
                                                      res_w=eval_params['res_w'],
                                                      axis=3,
                                                      phs_max=eval_params['phs_max'],
                                                      amp_max=None,
                                                      clamp=True,
                                                      normalize=True,
                                                      wavelength=self.hologram_params["wavelengths"])
        else:
            phs_only, amp_max = optics.tf_aadpm(cpx_out,
                                                propagator,
                                                depth_shift=eval_params["depth_shift"],
                                                adaptive_phs_shift=eval_params['adaptive_phs_shift'],
                                                batch=1,
                                                num_channels=3,
                                                res_h=eval_params['res_h'],
                                                res_w=eval_params['res_w'],
                                                sigma=eval_params['gaussian_sigma'],
                                                kernel_width=eval_params['gaussian_width'],
                                                phs_max=eval_params['phs_max'],
                                                amp_max=None,
                                                clamp=True,
                                                normalize=True,
                                                wavelength=self.hologram_params["wavelengths"])

        # restore pre-trained model
        self.saver = tf.compat.v1.train.Saver(max_to_keep=5)
        ckpt = tf.train.get_checkpoint_state(self.path_params["ckpt_parent_path"])
        if ckpt and ckpt.model_checkpoint_path:
            self.saver.restore(self.sess, tf.train.latest_checkpoint(self.path_params["ckpt_parent_path"]))
            print("model restored from %s" % self.path_params["ckpt_parent_path"])
        else:
            print("ERROR: NO RESTORED MODEL...")

        # load input
        rgb   = np.transpose(cv2.resize(cv2.imread(eval_params['rgb_path']),
            dsize=(eval_params['res_w'], eval_params['res_h']), interpolation=cv2.INTER_CUBIC), [2,0,1])[:3,:,:] / 255.0
        depth = cv2.resize(cv2.imread(eval_params['depth_path']),
            dsize=(eval_params['res_w'], eval_params['res_h']), interpolation=cv2.INTER_CUBIC) / 255.0
        if len(depth.shape) == 3:
            depth = depth[:,:,0]
        depth = depth[None,:,:]
        rgbd_val = np.concatenate((rgb, depth), axis=0)
        rgbd_val = rgbd_val[None,:,:,:].astype(np.float32)

        # evaluate output
        amp_out_val, phs_out_val, phs_only_val= self.sess.run([amp_out,
                                                               phs_out,
                                                               phs_only],
                                                               feed_dict={rgbd: rgbd_val})
        amp_out_val = np.transpose(amp_out_val[0,:,:,:], [1,2,0])
        phs_out_val = np.transpose(phs_out_val[0,:,:,:], [1,2,0])
        phs_only_val = np.transpose(phs_only_val[0,:,:,:], [1,2,0])

        # save output
        cv2.imwrite(os.path.join(eval_params["output_path"], "phs.png"), phs_out_val * 255.0)
        cv2.imwrite(os.path.join(eval_params["output_path"], "amp.png"), amp_out_val * 255.0)
        cv2.imwrite(os.path.join(eval_params["output_path"], "blue.png"), phs_only_val[:,:,0] * 255.0)
        cv2.imwrite(os.path.join(eval_params["output_path"], "green.png"), phs_only_val[:,:,1] * 255.0)
        cv2.imwrite(os.path.join(eval_params["output_path"], "red.png"), phs_only_val[:,:,2] * 255.0)

if __name__ == '__main__':
    cur_dir = os.path.dirname(os.path.realpath(__file__))

    parser = argparse.ArgumentParser()
    # evaluation parameters
    parser.add_argument('--eval-mode', action='store_true', help='Run in evaluation mode')
    parser.add_argument('--eval-res-h', default=1080, type=int, help='Input image height in evaluation mode')
    parser.add_argument('--eval-res-w', default=1920, type=int, help='Input image width in evaluation mode')
    parser.add_argument('--eval-rgb-path', default=os.path.join(cur_dir, "data", "example_input", "couch_rgb.png"), help='Input rgb image path in evaluation mode')
    parser.add_argument('--eval-depth-path', default=os.path.join(cur_dir, "data", "example_input", "couch_depth.png"), help='Input depth image path in evaluation mode')
    parser.add_argument('--eval-output-path', default=os.path.join(cur_dir, "data", "example_input"), help='Output directory for results')
    parser.add_argument('--eval-depth-shift', default=0, type=float, help='Depth shift (in mm) from the predicted midpoint hologram to the target hologram plane')
    parser.add_argument('--gaussian-sigma', default=0.7, type=float, help='Sigma of Gaussian kernel used by AA-DPM')
    parser.add_argument('--gaussian-width', default=3, type=int, help='Width of Gaussian kernel used by AA-DPM')
    parser.add_argument('--phs-max', default=3.0, type=float, help='Maximum phase modulation of SLM in unit of pi')
    parser.add_argument('--use-maimone-dpm', action='store_true', help='Use DPM of Maimone et al. 2017')

    # dataset parameters
    parser.add_argument('--dataset-res', default=384, type=int, help='dataset image resolution')
    parser.add_argument('--pitch', default=0.008, type=float, help='pixel pitch in mm')

    # model parameters
    parser.add_argument('--num-filters-per-layer', default=24, type=int, help='Number of filters per layer')
    parser.add_argument('--num-layers', default=30, type=int, help='Number layers')

    # validation parameters
    parser.add_argument('--validate-mode', action='store_true', help='Run in validation mode')

    opt = parser.parse_args()

    # fix random seed
    tf.compat.v1.set_random_seed(0)

    # hologram parameters, units in mm
    hologram_params = {
        "wavelengths" : np.array([0.000450, 0.000520, 0.000638]),  # laser wavelengths in BGR order
        "pitch" : opt.pitch,                                           # hologram pitch
        "res_h" : opt.dataset_res,                                 # dataset image height
        "res_w" : opt.dataset_res,                                 # dataset image width
        "depth_base"  : -3,                                        # input hologram plane (midpoint hologram)
        "depth_scale" : 6,                                         # 3D volume depth
        "double_pad": True                                         # double padding for propagation
    }

    # training parameters
    training_params = {
        "restore_trained_model": True,                             # flag to restore pre-trained model
        "batch" : 2,                                               # training batch
        "num_epochs": 1000,                                        # training epochs
        "decay_type": None,                                        # learning rate decay
        "decay_params": None,                                      # learning decay parameters
        "learning_rate" : 1e-4,                                    # learning rate
        "optimizer_type" : "adam",                                 # optimizer type
        "optimizer_params" :
            {"beta1":0.9, "beta2":0.99, "epsilon":1e-8},           # optimizer parameters
        "num_iter_per_model_save": 1000,                           # number of iterations per model save
        "num_iter_per_test": 1000,                                 # number of iterations per validation
        "num_top_depth_for_img_loss": 15,                          # number of top-k depths for computing focal stack loss
        "num_random_depth_for_img_loss": 5,                        # number of random depths selected from rest of the bins
        "depth_dependent_weight_scale": 0.35,                      # attention weight
        "num_hist_bins": 200                                       # number of focal stack bins
    }

    # model parameters
    model_params = {
        "name": "full_loss",                                       # model name
        "input_dim": 4,                                            # RGBD
        "output_dim": 6,                                           # amplitude+phase
        "num_layers": opt.num_layers,                              # number of convolution layers
        "interleave_rate": 1,                                      # interleaving rate (default to no interleaving)
        "num_filters_per_layer": opt.num_filters_per_layer,        # number of filters per convolution layer
        "filter_width": 3,                                         # filter width
        "bias_stddev": 0.01,                                       # bias standard deviation (for model initialization)
        "weight_var_scale": 0.25,                                  # weight variance (for model initialization)
        "renormalize_input": True,                                 # normalize input to [-0.5, 0.5]
        "activation_func": tf.nn.relu,                             # activation function for intermediate layers
        "output_activation_func": tf.nn.tanh,                      # activation function for the output layer
    }

    # loss function parameters
    use_l2_loss = False
    num_imgs_in_fs = training_params["num_top_depth_for_img_loss"] + training_params["num_random_depth_for_img_loss"]
    loss_params = {
        "use_l2_loss":   use_l2_loss,                              # activation function for the output layer
        "loss_op":       tf.losses.mean_squared_error if use_l2_loss else tf.compat.v1.losses.absolute_difference,
        "weight_holo":   1.0,                                      # hologram loss weight
        "weight_fs":     num_imgs_in_fs,                           # focal stack loss weight
        "weight_fs_tv":  num_imgs_in_fs,                           # focal stack tv loss weight
    }

    # path parameters
    labels = ["amp", "phs", "img", "depth"]
    train_base_path = os.path.join(cur_dir, "data", "train_%d" % opt.dataset_res)
    test_base_path = os.path.join(cur_dir, "data", "test_%d" % opt.dataset_res)
    validate_base_path = os.path.join(cur_dir, "data", "validate_%d" % opt.dataset_res)
    checkpoint_base_path = os.path.join(cur_dir, 'model', "ckpt_%s_pitch_%d_layers_%d_filters_%d" %
        (model_params["name"], hologram_params["pitch"]*1000, model_params["num_layers"], model_params["num_filters_per_layer"]))
    path_params = {
        "gen_record" : False,                                                           # generate tf record
        "labels" : labels,                                                              # image labels in the tfrecord
        "train_output_path"    : os.path.join(train_base_path, "train.tfrecord"),       # path to training set tfrecord
        "train_source_paths"   : [os.path.join(train_base_path, x) for x in labels],    # path to training set raw images
        "test_output_path"     : os.path.join(test_base_path, "test.tfrecord"),         # path to test set tfrecord
        "test_source_paths"    : [os.path.join(test_base_path, x) for x in labels],     # path to test set raw images
        "validate_output_path" : os.path.join(validate_base_path, "validate.tfrecord"), # path to validate set tfrecord
        "validate_source_paths": [os.path.join(validate_base_path, x) for x in labels], # path to validation set raw images
        "ckpt_path"            : os.path.join(checkpoint_base_path, "ckpt"),            # checkpoint path
        "ckpt_parent_path"     : checkpoint_base_path,                                  # checkpoint parent path
        "epoch_record_file"    : "epoch_record.txt",                                    # txt file that records how many epochs have been trained
        "log_path"             : "logs",                                                # log file
    }

    # training dataset parameters
    train_dataset_params = {
        "repeat": True,
        "sample_count": 3800,
        "batch": training_params["batch"],
        "res_h": hologram_params["res_h"],
        "res_w": hologram_params["res_w"],
        "prefetch_buffer_size": 2,
        "num_parallel_calls": 4,
        "shuffle_buffer_size": 2,
        "num_epochs": training_params["num_epochs"]
    }

    # test dataset parameters
    test_dataset_params = {
        "repeat": True,
        "sample_count": 100,
        "batch": training_params["batch"],
        "res_h": hologram_params["res_h"],
        "res_w": hologram_params["res_w"],
        "num_parallel_calls": 2,
        "prefetch_buffer_size": 4,
        "shuffle_buffer_size": 2,
        "num_epochs": training_params["num_epochs"]
    }

    # validation dataset parameters
    validate_dataset_params = {
        "repeat": True,
        "sample_count": 100,
        "batch": training_params["batch"],
        "res_h": hologram_params["res_h"],
        "res_w": hologram_params["res_w"],
        "num_parallel_calls": 2,
        "prefetch_buffer_size": 4,
        "shuffle_buffer_size": 2,
        "num_epochs": training_params["num_epochs"]
    }

    # build model
    tensor_holo_model = TensorHolographyModel(hologram_params=hologram_params,
                                              training_params=training_params,
                                              model_params=model_params,
                                              loss_params=loss_params,
                                              path_params=path_params,
                                              train_dataset_params=train_dataset_params,
                                              test_dataset_params=test_dataset_params,
                                              validate_dataset_params=validate_dataset_params)

    # train model from scratch (will be available in the second release)
    # tensor_holo_model.train()

    # validate pre-trained model
    if opt.validate_mode:
        tensor_holo_model.validate()

    # evaluate pre-trained model
    if opt.eval_mode:
        eval_params = {
            "res_h": opt.eval_res_h,
            "res_w": opt.eval_res_w,
            "rgb_path": opt.eval_rgb_path,
            "depth_path": opt.eval_depth_path,
            "output_path": opt.eval_output_path,
            "double_pad": True,
            "use_maimone_dpm": opt.use_maimone_dpm,
            "adaptive_phs_shift": False,
            "depth_shift": opt.eval_depth_shift,
            "gaussian_sigma": opt.gaussian_sigma,
            "gaussian_width": opt.gaussian_width,
            "phs_max": [opt.phs_max*np.pi] * 3,
            "amp_max": None,
        }
        tensor_holo_model.evaluate(eval_params)


import paddle
import itertools
import log
import sys
import time
import datetime
import paddle.vision.transforms as T
from paddle.vision.datasets import Cifar10
from paddle.optimizer import Adam
from paddle.optimizer.lr import CosineAnnealingDecay
from paddle.optimizer.lr import LinearWarmup
from paddle.regularizer import L2Decay
from paddle.nn import CrossEntropyLoss
from paddle.metric import Accuracy
from Edge_blur import EdgeBlurs
import ViT.module

class Cosine(CosineAnnealingDecay):
    def __init__(self, lr, step_each_epoch, epochs, **kwargs):
        super(Cosine, self).__init__(learning_rate=lr, T_max=step_each_epoch * epochs)
        self.update_specified = False

class CosineWarmup(LinearWarmup):
    def __init__(self, lr, step_each_epoch, epochs, warmup_epoch=5, **kwargs):
        assert epochs > warmup_epoch, "total epoch({}) should be larger than warmup_epoch({}) in CosineWarmup.".format(
            epochs, warmup_epoch)
        warmup_step = warmup_epoch * step_each_epoch
        start_lr = 0.0
        end_lr = lr
        lr_sch = Cosine(lr, step_each_epoch, epochs - warmup_epoch)

        super(CosineWarmup, self).__init__(learning_rate=lr_sch, warmup_steps=warmup_step, start_lr=start_lr,
                                           end_lr=end_lr)
        self.update_specified = False

sys.stdout = log.Logger('./log/log{}.txt'.format(
    time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime(time.time()))),
    mode='w', encoding='utf-8')

if __name__ == '__main__':
    start_time = datetime.datetime.now()
    print(f'训练开始，当前时间为{start_time}')

    print("=======begin=======")

    img_size = 128

    drop = 0.2

    net = ViT.module.VisionTransformer(
        img_size=img_size,
        patch_size=16,
        class_dim=10,
        embed_dim=384,
        depth=4,
        num_heads=6,
        drop_rate=drop,
        attn_drop_rate=drop,
        drop_path_rate=drop,
        mlp_ratio=4,
        qkv_bias=True,
        epsilon=1e-6)
    model = paddle.Model(net)
    model.load(r'pre param/63.pdparams')
    edge_blurs = EdgeBlurs()
    train_transform = T.Compose(
        [
            T.Resize(img_size),
            edge_blurs,
            T.RandomErasing(),
            T.RandomHorizontalFlip(),
            T.RandomVerticalFlip(),
            T.ToTensor(),
            T.Normalize(
                mean=[0.5, 0.5, 0.5],
                std=[0.5, 0.5, 0.5],
                to_rgb=True,
            ),
        ]
    )

    test_transform = T.Compose(
        [
            T.Resize(img_size),
            T.ToTensor(),
            T.Normalize(
                mean=[0.5, 0.5, 0.5],
                std=[0.5, 0.5, 0.5],
                to_rgb=True,
            ),
        ]
    )

    train_dataset = Cifar10(mode='train', transform=train_transform, download=True, backend='cv2')
    # print(len(train_dataset))
    test_dataset = Cifar10(mode='test', transform=test_transform, download=True, backend='cv2')
    # print(len(test_dataset))

    epochs = 300
    warmup_epoch = 10

    # for img, label in itertools.islice(iter(train_dataset), 5):
    #     print(type(img), img.shape, label)

    lr = 1e-3
    # scheduler = CosineWarmup(lr=lr, step_each_epoch=100, epochs=epochs, warmup_epoch=warmup_epoch, start_lr = 0, end_lr = lr, verbose = True)
    # scheduler = LinearWarmup(learning_rate=lr, warmup_steps=warmup_epoch, start_lr=0, end_lr=lr)
    # scheduler = CosineAnnealingDecay(learning_rate=lr, T_max=warmup_epoch, eta_min=0)
    optimizer = Adam(learning_rate=lr, parameters=model.parameters())
    model.prepare(optimizer, CrossEntropyLoss(), Accuracy(topk=(1, 5)))

    print()
    print(f'the lr is {lr}')
    print(f'the epochs is {epochs}, and the warmup epoch is {warmup_epoch}')
    print()

    use_gpu = True
    paddle.set_device('gpu:0')
    model.fit(train_dataset,
              test_dataset,
              epochs=epochs,
              batch_size=300,
              save_dir="param",
              save_freq=7,
              log_freq=1,
              num_workers=0,
              shuffle=True,
              eval_freq=7,
              verbose=1)

    print("========end========")

    end_time = datetime.datetime.now()
    print(f'训练结束，当前时间为{end_time}')


import os
import time
import sys
import traceback
import glob

import cv2
import numpy as np
import onnxruntime

def main(args):
    classes = ['nsfw', 'sfw']

    session =  onnxruntime.InferenceSession(args.model_path, providers=['CUDAExecutionProvider'])
    input_name = session.get_inputs()[0].name
    output_name = [output.name for output in session.get_outputs()]

    tp, fp, tn, fn = 0, 0, 0, 0
    total_num = 0
    total_used_time = 0.0
    for ipath in glob.glob("/data/test/*/*.jpg", recursive=True):
        image = cv2.imread(ipath)
        if image is not None:
            total_num += 1
            label = ipath.split('/')[-2]
            if total_num >= 5:
                bg_time = time.time()
            blob = cv2.dnn.blobFromImage(image, 1.0/127.5, (224, 224), (127.5, 127.5, 127.5), swapRB=True)
            out = session.run(output_name, {input_name: blob})[0][0]
            if total_num >= 5:
                total_used_time += (time.time() - bg_time)
            pred = classes[np.argmax(out)]
            if pred == label == "positive":
                tp += 1
            elif pred == label == "negative":
                tn += 1
            elif pred == "negative" and label == "positive":
                fn += 1
            else:
                fp += 1
    print('Total test number is: {}, average used time is: {}'.format(total_num, total_used_time/(total_num-4)))
    """
    FPR: False Positive Rate,  FPR = FP / N = FP / (FP+TN) = 1 - TNR
    FNR: False Negtive Rate, FNR = FN / P = FN / (FN+TP) = 1 - TPR
    TPR: True Positive Rate, TPR = TP / p = TP / (TP+FN) = 1 - FNR
    TNR: True Negtive Rate, TNR = TN / N = TN / (TN+FP) = 1 - FPR
    Accuracy = (TP+TN) / (TP+FP+FN+TN)
    Precision(PPV) = TP / (TP+FP)
    Recall(TPR) = TP / (TP+FN) = 1 - FNR
    F1 score = 2 * precision * recall / (precision + recall)
    要求FPR和FNR越低越好
    """

    fpr = fp / (fp + tn)
    fnr = fn / (fn + tp)
    accuracy = (tp + tn) / (tp + fp + fn + tn)
    precision = tp / (tp + fp)
    recall = tp / (tp + fn)
    f1score = 2 * precision * recall / (precision + recall)

    print('FPR={}, FNR={}, accuracy={}, precision={}, recall={}, f1score={}'.format
          (fpr, fnr, accuracy, precision, recall, f1score))

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description='ONNX GPU Inference')

    parser.add_argument('--model-path', default='./models/model.onnx', help='onnx model')
    parser.add_argument('--resize', default=224, type=int, help='size of resize')

    args = parser.parse_args()

    print(args)
    main(args)


from common.lettuce_utils import *
import cssselect
from lettuce import step
from nose.tools import assert_equals, assert_not_equals, assert_raises
from user.models import TinvilleUser
import lettuce.django
from selenium.common.exceptions import *

@step(u'Given the beta shop')
def given_the_beta_shop(step):
    world.browser.get(lettuce.django.get_server().url('/BetaShop'))
    wait_for_browser_to_have_url(lettuce.django.get_server().url('/access_code?shop=BetaShop'))

@step(u'Given a subsequent visit to the beta shop')
def given_a_subsequent_visit_to_the_beta_shop(step):
    world.browser.get(lettuce.django.get_server().url('/BetaShop'))

@step(u'Then the access code page is displayed')
def then_the_access_code_page_is_displayed(step):
    assert_id_exists('betaform')

@step(u'And I enter the access code')
def and_i_enter_the_access_code(step):
    user = TinvilleUser.objects.get(email="joe@schmoe.com")
    form = fill_in_access_form(access_code=user.access_code)
    form.submit()

@step(u'Then I should be redirected to the beta shop')
def then_I_should_be_redirected_to_the_beta_shop(step):
    world.browser.get(lettuce.django.get_server().url('/BetaShop'))


import re
import ast

from app import redis_store

def new_cache(storage_type, *args, **kwargs):
    # storage_type 为储存类型 'list','detail','cache'
    def decorator(func):
        def wrapper(*args, **kwargs):
            args_data = [data for data in args]
            #args_data = ['origin','faculty','page','force_reload']
            '''
            若是获取LIST，参数为 origin,faculty,page,force_reload
            若是获取DETAIL，参数为 url,request_type

            '''
            if storage_type == 'list':
                # 缓存新闻列表
                name = args_data[0]
                # origin
                if name in ['xy', 'jw']:
                    # 将xy,jw与xb区分,方便进行缓存
                    key = args_data[2]
                    # page
                else:
                    key = "%s_%s" % (args_data[1], args_data[2])
                    # faculty_page
                data = redis_store.hget(name, key)

                if args_data[3]:
                    data = func(*args, **kwargs)
                    redis_store.hset(name, key, str(data))
                    redis_store.expire(name, 2592000)
                    # 缓存过期时间为30天
                elif data:
                    data = ast.literal_eval(bytes.decode(data))
                    redis_store.expire(name, 2592000)
                    return data

                else:
                    return {}
                '''
                else:
                    data = func(*args, **kwargs)
                    redis_store.hset(name, key, str(data))
                    redis_store.expire(name, 2592000)
                    # 缓存过期时间为30天
                '''

            elif storage_type == 'detail':
                # 缓存新闻详细
                url = args[0]
                list_length = len(args)
                # url
                result_1 = re.search('http://(.*?)/', url)
                # 用于判断是是否是教务处新闻详细
                if result_1.group(1) == 'www.gdust.cn':
                    result_2 = re.search('(\d{8})\/(\d{1,4})', url)
                    kw = result_2.group(1) + result_2.group(2)
                else:
                    kw = re.search('\&id\=(\d{1,3})', url).group(1)

                data = redis_store.get(kw)
                if args[1]:
                    data = func(*args, **kwargs)
                    redis_store.set(kw, str(data))
                    redis_store.expire(kw, 2592000)
                elif data:
                    data = ast.literal_eval(bytes.decode(data))
                    redis_store.expire(kw, 2592000)
                    return data
                else:
                    return {}
                '''
                else:
                    data = func(*args, **kwargs)
                    redis_store.set(kw, str(data))
                    redis_store.expire(kw, 2592000)
                '''

            return data
        return wrapper
    return decorator


"""Added double indexes

Revision ID: 5783809706d2
Revises: 4e2dde2d15e6
Create Date: 2015-06-01 21:08:15.106408

"""

# revision identifiers, used by Alembic.
revision = '5783809706d2'
down_revision = '4e2dde2d15e6'

from alembic import op
import sqlalchemy as sa

def upgrade():
    ### commands auto generated by Alembic - please adjust! ###
    op.create_index('ix_ecc_pi_ecc', 'ecc', ['pi', 'ecc'], unique=False)
    op.create_index('ix_station_spid_radioepg_enabled', 'station', ['service_provider_id', 'radioepg_enabled'], unique=False)
    ### end Alembic commands ###

def downgrade():
    ### commands auto generated by Alembic - please adjust! ###
    op.drop_index('ix_station_spid_radioepg_enabled', table_name='station')
    op.drop_index('ix_ecc_pi_ecc', table_name='ecc')
    ### end Alembic commands ###


from funciones import inicio, dado, verificarPosiciones, verificarCasilla, flujoJuego

#Me permite probar que el número que me proporciona la función dado( ) si esta en el rango
#EL dado va de 1 a 6

def test_dado_aleatorio():
    #este ciclo es para hacerlo un par de veces.
    for x in range(3):
        randomNumber  = dado()
        assert randomNumber>0 and randomNumber<6


"""bitforum URL Configuration

The `urlpatterns` list routes URLs to views. For more information please see:
    https://docs.djangoproject.com/en/3.1/topics/http/urls/
Examples:
Function views
    1. Add an import:  from my_app import views
    2. Add a URL to urlpatterns:  path('', views.home, name='home')
Class-based views
    1. Add an import:  from other_app.views import Home
    2. Add a URL to urlpatterns:  path('', Home.as_view(), name='home')
Including another URLconf
    1. Import the include() function: from django.urls import include, path
    2. Add a URL to urlpatterns:  path('blog/', include('blog.urls'))
"""
from django.contrib import admin
from django.urls import path
from django.urls import path, include
from posts import urls
from .views import *
from django.conf import settings
from django.conf.urls.static import static

urlpatterns = [
    path('admin/', admin.site.urls),
    # path('admin',admin.site.urls),
    path('myProfile',myProfile),
    path('login',login_signup_page),
    path('checkLogin',checkLogin,name='checkLogin'),
    path('',home),
    path('logout', logout),
    path('deleteFollowing', deleteFollowing,name='deleteFollowing'),
    path('unfollowTopic', unfollowTopic,name='unfollowTopic'),
    path('removeFollower', removeFollower, name='removeFollower'),
    path('post/<int:pid>', viewPost, name='viewPost'),
    path('signup', signup,name='signup'),
    path('viewProfile/<int:uid>',viewProfile,name='viewProfile'),
    path('doUpvote',doUpvote,name='doUpvote'),
    path('doDownvote',doDownvote,name='doDownvote'),
    path('followAjax',followajax,name='followAjax'),
    path('followTopicAjax',followTopicAjax,name='followTopicAjax'),
    path('unfollowUserBtnAjax',unfollowUserBtnAjax,name='unfollowUserBtnAjax'),
    path('followUserBtnAjax',followUserBtnAjax,name='followUserBtnAjax'),
    path('postComment/<int:pid>',postComment,name='postComment'),
    path('downComment',downComment,name='downComment'),
    path('upComment',upComment,name='upComment'),
    path('makePost',makePost,name='makePost'),
    path('fetchNotification',fetchNotification,name='fetchNotification'),
    path('mark_all_as_read',mark_all_as_read,name='mark_all_as_read'),
    path('deletePost',deletePost,name='deletePost'),
    path('saveEditProfile',saveEditProfile,name='saveEditProfile'),
    path('search',search,name='search')

]

urlpatterns += static(settings.MEDIA_URL,document_root=settings.MEDIA_ROOT)


import asyncio
from typing import Optional, Any, List, Dict, Union

from simnet.client.components.chronicle.base import BaseChronicleClient
from simnet.client.routes import RECORD_URL
from simnet.errors import DataNotPublic, BadRequest
from simnet.models.genshin.chronicle.abyss import SpiralAbyss, SpiralAbyssPair
from simnet.models.genshin.chronicle.characters import Character
from simnet.models.genshin.chronicle.notes import Notes, NotesWidget, NotesOverseaWidget
from simnet.models.genshin.chronicle.stats import (
    PartialGenshinUserStats,
    GenshinUserStats,
    FullGenshinUserStats,
)
from simnet.models.lab.record import RecordCard
from simnet.utils.enum_ import Game, Region
from simnet.utils.player import recognize_genshin_server, recognize_region

__all__ = ("GenshinBattleChronicleClient",)

class GenshinBattleChronicleClient(BaseChronicleClient):
    """A client for retrieving data from Genshin's battle chronicle component.

    This class is used to retrieve various data objects from StarRail's battle chronicle component,
    including real-time notes, user statistics, and character information.
    """

    async def _request_genshin_record(
        self,
        endpoint: str,
        player_id: Optional[int] = None,
        method: str = "GET",
        endpoint_type: str = "api",
        lang: Optional[str] = None,
        payload: Optional[Dict[str, Any]] = None,
    ):
        """Get an arbitrary object from StarRail's battle chronicle.

        Args:
            endpoint (str): The endpoint of the object to retrieve.
            endpoint_type (str, optional): The type of endpoint to send the request to.
            player_id (Optional[int], optional): The player ID. Defaults to None.
            method (str, optional): The HTTP method to use. Defaults to "GET".
            lang (Optional[str], optional): The language of the data. Defaults to None.
            payload (Optional[Dict[str, Any]], optional): The request payload. Defaults to None.

        Returns:
            Dict[str, Any]: The requested object.

        Raises:
            BadRequest: If the request is invalid.
            DataNotPublic: If the requested data is not public.
            ValueError: If the player ID is not specified.
        """
        player_id = player_id or self.player_id

        if player_id is None:
            raise ValueError("Player ID is not specified.")

        if payload is None:
            payload = dict(role_id=player_id, server=recognize_genshin_server(player_id))
        else:
            payload = dict(role_id=player_id, server=recognize_genshin_server(player_id), **payload)

        data, params = None, None
        if method == "POST":
            data = payload
        else:
            params = payload

        return await self.request_game_record(
            endpoint,
            endpoint_type=endpoint_type,
            lang=lang,
            game=Game.GENSHIN,
            region=recognize_region(player_id, game=Game.GENSHIN),
            params=params,
            data=data,
        )

    async def get_partial_genshin_user(
        self,
        player_id: Optional[int] = None,
        *,
        lang: Optional[str] = None,
    ) -> PartialGenshinUserStats:
        """Get partial genshin user without character equipment.

        Args:
            player_id (Optional[int], optional): The player ID. Defaults to None.
            lang (Optional[str], optional): The language of the data. Defaults to None.

        Returns:
            PartialGenshinUserStats: The requested partial genshin user without character equipment.
        """
        data = await self._request_genshin_record("index", player_id, lang=lang)
        return PartialGenshinUserStats(**data)

    async def get_genshin_characters(
        self,
        player_id: Optional[int] = None,
        *,
        lang: Optional[str] = None,
    ) -> List[Character]:
        """Get genshin user characters.

        Args:
            player_id (Optional[int], optional): The player ID. Defaults to None.
            lang (Optional[str], optional): The language of the data. Defaults to None.

        Returns:
            Character: The requested genshin user characters.
        """
        data = await self._request_genshin_record("character", player_id, lang=lang, method="POST")
        return [Character(**i) for i in data["avatars"]]

    async def get_genshin_user(
        self,
        player_id: Optional[int] = None,
        *,
        lang: Optional[str] = None,
    ) -> GenshinUserStats:
        """Get genshin user stats.

        Args:
            player_id (Optional[int], optional): The player ID. Defaults to None.
            lang (Optional[str], optional): The language of the data. Defaults to None.

        Returns:
            GenshinUserStats: The requested genshin user stats.
        """
        data, character_data = await asyncio.gather(
            self._request_genshin_record("index", player_id, lang=lang),
            self._request_genshin_record("character", player_id, lang=lang, method="POST"),
        )
        data = {**data, **character_data}

        return GenshinUserStats(**data)

    async def get_genshin_spiral_abyss(
        self,
        player_id: Optional[int] = None,
        *,
        previous: bool = False,
        lang: Optional[str] = None,
    ) -> SpiralAbyss:
        """Get genshin spiral abyss runs.

        Args:
            player_id (Optional[int], optional): The player ID. Defaults to None.
            previous (bool, optional): Whether to retrieve the data for the previous season of the Spiral Abyss.
                Defaults to False.
            lang (Optional[str], optional): The language of the data. Defaults to None.

        Returns:
            SpiralAbyss: genshin spiral abyss runs.
        """
        payload = dict(schedule_type=2 if previous else 1)
        data = await self._request_genshin_record("spiralAbyss", player_id, lang=lang, payload=payload)

        return SpiralAbyss(**data)

    async def get_genshin_notes(
        self,
        player_id: Optional[int] = None,
        *,
        lang: Optional[str] = None,
        autoauth: bool = True,
    ) -> Notes:
        """Get Genshin's real-time notes.

        Args:
            player_id (Optional[int], optional): The player ID. Defaults to None.
            lang (Optional[str], optional): The language of the data. Defaults to None.
            autoauth (bool, optional): Whether to automatically authenticate the user. Defaults to True.

        Returns:
            Notes: The requested real-time notes.

        Raises:
            BadRequest: If the request is invalid.
            DataNotPublic: If the requested data is not public.
        """
        try:
            data = await self._request_genshin_record("dailyNote", player_id, lang=lang)
        except DataNotPublic as e:
            # error raised only when real-time notes are not enabled
            if player_id and self.player_id != player_id:
                raise BadRequest(e.response, "Cannot view real-time notes of other users.") from e
            if not autoauth:
                raise BadRequest(e.response, "Real-time notes are not enabled.") from e

            await self.update_settings(3, True, game=Game.GENSHIN)
            data = await self._request_genshin_record("dailyNote", lang=lang)

        return Notes(**data)

    async def get_full_genshin_user(
        self,
        player_id: Optional[int] = None,
        *,
        lang: Optional[str] = None,
    ) -> FullGenshinUserStats:
        """Get a genshin user with all their possible data.

        Args:
            player_id (Optional[int], optional): The player ID. Defaults to None.
            lang (Optional[str], optional): The language of the data. Defaults to None.

        Returns:
            Character: The requested genshin user with all their possible data.
        """
        index, character, abyss1, abyss2, activities = await asyncio.gather(
            self._request_genshin_record("index", player_id, lang=lang),
            self._request_genshin_record("character", player_id, lang=lang, method="POST"),
            self.get_genshin_spiral_abyss(player_id, lang=lang, previous=False),
            self.get_genshin_spiral_abyss(player_id, lang=lang, previous=True),
            self.get_genshin_activities(player_id, lang=lang),
        )
        user = {**index, **character}
        abyss = SpiralAbyssPair(current=abyss1, previous=abyss2)

        return FullGenshinUserStats(**user, abyss=abyss, activities=activities)

    async def get_genshin_activities(self, player_id: Optional[int] = None, *, lang: Optional[str] = None) -> Dict:
        """Get genshin activities.

        Args:
            player_id (Optional[int], optional): The player ID. Defaults to None.
            lang (Optional[str], optional): The language of the data. Defaults to None.

        Returns:
            Dict: The requested get_genshin_activities.
        """
        return await self._request_genshin_record("activities", player_id, lang=lang)

    async def get_record_card(
        self,
        account_id: Optional[int] = None,
        *,
        lang: Optional[str] = None,
    ) -> Optional[RecordCard]:
        """Get a genshin player record cards.

        Args:
            account_id: Optional[int], the user's account ID, defaults to None
            lang: Optional[str], the language version of the request, defaults to None

        Returns:
            Genshin player record cards.

        Returns:
            Optional[RecordCard]: RecordCard objects.
        """
        account_id = account_id or self.account_id

        record_cards = await self.get_record_cards(account_id, lang=lang)

        for record_card in record_cards:
            if record_card.game == Game.GENSHIN:
                return record_card

        return None

    async def get_genshin_notes_by_stoken(
        self,
        lang: Optional[str] = None,
    ) -> Union[NotesWidget, NotesOverseaWidget]:
        """Get Genshin's real-time notes.

        Args:
            lang (Optional[str], optional): The language of the data. Defaults to None.

        Returns:
            NotesWidget (Union[NotesWidget, NotesOverseaWidget]): The requested real-time notes.

        Raises:
            BadRequest: If the request is invalid.
        """
        stoken = self.cookies.get("stoken")
        if stoken is None:
            raise ValueError("stoken not found in cookies.")
        stuid = self.cookies.get("stuid")
        if stuid is None and self.account_id is None:
            raise ValueError("account_id or stuid not found")
        if self.account_id is not None and stuid is None:
            self.cookies.set("stuid", str(self.account_id))
        if self.region == Region.OVERSEAS:
            route = RECORD_URL.get_url(self.region) / "../community/apihub/api/widget/data"
            params = {"game_id": "2"}
            data = await self.request_lab(route, params=params, lang=lang)
            model = NotesOverseaWidget
        else:
            data = await self._request_genshin_record("widget/v2", endpoint_type="aapi", lang=lang)
            model = NotesWidget
        return model(**data)


#!/usr/bin/env python

import sys
import os

import numpy
import scipy.interpolate as interp
import matplotlib.pyplot as plt

import clawpack.visclaw.colormaps as colormaps

def read_deformation(path):
    r""""""

    data = numpy.loadtxt(path)

    x = data[:,0]
    y = data[:,1]
    z = data[:,2]

    # Find maximum extents and minimum distance between points
    extents = [numpy.min(x), numpy.min(y), numpy.max(x), numpy.max(y)]

    return extents, x, y, z

def project_deformation(extents, x, y, z):
    r""""""

    # Construct new grid
    num_cells = [ int((extents[2] - extents[0]) / numpy.min(numpy.abs(numpy.diff(x)))),
                  int((extents[3] - extents[1]) / numpy.min(numpy.abs(numpy.diff(y)))) ]
    t = numpy.array([0.0,1.0])
    X, Y = numpy.meshgrid(numpy.linspace(extents[0],extents[2],num_cells[0]),
                          numpy.linspace(extents[1],extents[3],num_cells[1]))
    Z = numpy.zeros((num_cells[1], num_cells[0], 2))
    Z[:,:,1] = interp.griddata((x, y), z, (X, Y), method='linear', fill_value=0.0)

    return t, X, Y, Z

def write_deformation_to_file(t, X, Y, Z, output_file, topo_type=1):
    r"""Write out a dtopo file to *output_file*

    input
    -----
     - *t* (numpy.ndarray[:]) - Array containing time points, note that
       *t.shape[0] == Z.shape[2]*.
     - *X* (numpy.ndarray[:,:]) - Array containing x-coodinates (longitude),
       should be in the form given by *numpy.meshgrid*.
     - *Y* (numpy.ndarray[:,:]) - Array containing y-coordinates (latitude),
       should be in the form given by *numpy.meshgrid*.
     - *Z* (numpy.ndarray[:,:,:]) - Array containing deformation from original
       bathymetry.
     - *output_file* (path) - Path to the output file to written to.
     - *topo_type* (int) - Type of topography file to write out.  Default is 1.

    """

    # Temporary catch for non implemented topo_type input
    if topo_type != 1:
        raise NotImplementedError("Topography types 2 and 3 are not yet supported.")

    # Construct each interpolating function and evaluate at new grid
    try:
        outfile = open(output_file, 'w')

        if topo_type == 1:
            # Topography file with 4 columns, t, x, y, dz written from the upper
            # left corner of the region
            Y_flipped = numpy.flipud(Y)
            for n in xrange(t.shape[0]):
                Z_flipped = numpy.flipud(Z[:,:,n])
                for j in xrange(Y.shape[0]):
                    for i in xrange(X.shape[1]):
                        outfile.write("%s %s %s %s\n" % (t[n], X[j,i], Y_flipped[j,i], Z_flipped[j,i]))

        elif topo_type == 2 or topo_type == 3:
            raise NotImplementedError("Topography types 2 and 3 are not yet supported.")
        else:
            raise ValueError("Only topography types 1, 2, and 3 are supported.")

    except IOError as e:
        raise e
    finally:
        outfile.close()

def read_dtopo_file(path, topo_type=1):
    r""""""

    if topo_type != 1:
        raise ValueError("Topography type != 1 is not implemented.")

    # Load raw data
    data = numpy.loadtxt(path)

    # Parse data
    t = data[:,0]
    x = data[:,1]
    y = data[:,2]

    # Initialize extents
    t0 = t[0]
    lower = [x[0], y[0]]
    upper = [None, None]
    num_cells = [0,0]

    # Count total x-values
    for row in xrange(1,data.shape[0]):
        if x[row] == x[0]:
            num_cells[0] = row
            break

    # Count total y-values
    for row in xrange(num_cells[0], data.shape[0]):
        if t[row] != t0:
            num_cells[1] = row / num_cells[0]
            num_times = data.shape[0] / row
            break

    # Check extents
    assert(t[0] != t[num_cells[0] * num_cells[1] + 1])
    # assert(t[0] == t[num_cells[0] * num_cells[1]])

    # Fill in rest of pertinent data
    t = data[::num_cells[0] * num_cells[1], 0]
    x = data[:num_cells[0], 1]
    y = data[:num_cells[0] * num_cells[1]:num_cells[0], 2]
    upper = [x[-1], y[-1]]
    X, Y = numpy.meshgrid(x, y)
    Z = numpy.empty( (num_times, num_cells[0], num_cells[1]) )

    for (n,time) in enumerate(t):
        Z[n,:,:] = data[num_cells[0] * num_cells[1] * n:
                        num_cells[0] * num_cells[1] * (n+1), 3].reshape(
                                (num_cells[0], num_cells[1]))

    return t, X, Y, Z

if __name__ == "__main__":

    # path = os.path.abspath("../tohoku2011-paper1/sources/Ammon.txydz")
    path = os.path.abspath("./initFig3a.dat")

    if len(sys.argv) > 1:
        path = sys.argv[1]

    extents, x, y, z = read_deformation(path)

    # Plot data
    colorbar_limits = [numpy.min(z), numpy.max(z)]
    if -colorbar_limits[0] >= colorbar_limits[1]:
        colorbar_limits[1] = -colorbar_limits[0]
    else:
        colorbar_limits[0] = -colorbar_limits[1]
    cmap = colormaps.make_colormap({1.0:'r',0.5:'w',0.0:'b'})

    # Plot original data
    fig, axes = plt.subplots(1)
    plot = axes.scatter(x, y, c=z, vmin=colorbar_limits[0],
                                   vmax=colorbar_limits[1], cmap=cmap)
    axes.set_title("Original Data")
    axes.set_aspect('equal')
    cb = fig.colorbar(plot, ax=axes)
    cb.set_label('Deformation (m)')

    # # Project and write out deformation
    t, X, Y, Z = project_deformation(extents, x, y, z)
    output_file = "./saito.xyzt"
    write_deformation_to_file(t, X, Y, Z, output_file)

    # Plot new data
    fig, axes = plt.subplots(2)
    import pdb; pdb.set_trace()
    for frame in xrange(2):
        im = axes[frame].pcolormesh(X, Y, Z[:,:,frame], vmin=colorbar_limits[0],
                                                        vmax=colorbar_limits[1],
                                                        cmap=cmap)
        axes[frame].set_title("Projected Data")
        axes[frame].set_aspect('equal')
    cb = fig.colorbar(im, ax=axes[1])
    cb.set_label('Deformation (m)')

    plt.show()