from functools import reduce

import logging
import pickle
import multiprocessing as mp
from multiprocessing import Pool
import numpy as np
import numba as nb
import pde
import cantera as ct
from scipy import optimize

class CombustionChamber:
    def __init__ (self, mdot, ct_object, burned_state, hA=700):
        self.mdot = mdot                        # kg/s
        self.gas = ct_object                    # gas object
        self.eq_state = burned_state            # HP equilibrium state
        self.gas.TPX = burned_state             # Set equilibrium state
        T0, P0, X0 = self.gas.TPX
        self.T0 = T0                            # K, adiabatic flame temperature
        self.P0 = P0                            # Pa, pressure
        self.X0 = X0                            # Composition in mole fractions, Fuel + Air
        self.h0 = self.gas.enthalpy_mass        # inlet enthalpy
        self.hA = hA                            # HTC x Area
        self.T1 = T0
        self.Twall0 = 1100 + 273.15
        self.Twall1 = 1100 + 273.15
        self.Area = 6.7 * 16.7

    def update_mdot (self, mdot_new):
        if mdot_new : self.mdot = mdot_new

    def update_Twall (self, Twall0=None, Twall1=None):
        if Twall0: self.Twall0 = Twall0
        if Twall1: self.Twall1 = Twall1

    def energy_balance_equation (self, Tout):
        self.gas.TP = Tout, None
        h1 = self.gas.enthalpy_mass
        q1, q2 = self.heat(Tout)

        return (self.mdot * (self.h0 - h1) - q1 - q2)

    def solve (self, ):
        """ Iteratively solve for outlet temperature that balance with heat loss to walls """

        meanTwall = (self.Twall0 + self.Twall1) / 2
        T_low = meanTwall - (self.T0 - meanTwall)
        try:
            f_found = optimize.root_scalar(self.energy_balance_equation,
                                        bracket=[T_low, self.T0])
            self.T1 = f_found.root
        except ValueError:
            self.T1 = meanTwall

        return f_found.root

    def heat (self, Tout=None):
        ''' Heat(W) to walls '''
        if Tout is None: Tout = self.T1
        Tgas = (self.T0 + Tout) / 2
        
        return self.hA * (Tgas - self.Twall0), self.hA * (Tgas - self.Twall1)

class CokeCharge:
    
    def __init__ (self, t_charge, idx_oven):
        self.t_charge = t_charge
        self.t_push = None
        self.idx_oven = idx_oven
        self.Q = 0

    def bake (self, dQ):
        self.Q += dQ
        
    def end_baking (self, t):
        self.t_push = t

brick_thickness = 0.14                            # m, 
n_grid_brick = 16                                 # Number of Grid points inside
wall_grid = pde.CartesianGrid([[0, brick_thickness]], n_grid_brick, periodic=False)
wall_area = 6.7 * 16.7 # m^2 , Oven cross section area

# op_grad2 = wall_grid.make_operator_no_bc('gradient_squared', backend='scipy')
# op_grad = wall_grid.make_operator_no_bc('gradient', backend='scipy')
# op_lap = wall_grid.make_operator_no_bc('laplace', backend='scipy')

# op_info_grad2 = wall_grid._get_operator_info('gradient_squared')
# op_info_grad = wall_grid._get_operator_info('gradient')
# op_info_lap = wall_grid._get_operator_info('laplace')

class CokeOvenBrickHeatEqn(pde.PDEBase):
    """Implementation of the Heat equation"""

    def __init__ (self, bc="auto_periodic_neumann"):
        self.bc = bc
        self.rho = 1900            # kg / m3
        self.kCoef0 = 0.93         # W / m / K
        self.kCoef1 = 0.698e-3     # W / m / K2
        self.cpCoef0 = 837.2       # J / kg / K
        self.cpCoef1 = 251.2e-3    # J / kg / K2

    def k (self, T):
        return T * self.kCoef1 + self.kCoef0
    
    def cp (self, T):
        return T * self.cpCoef1 + self.cpCoef0
        
    def update_bc (self, gradT_chamber=None, T_oven=None):
        bc0, bc1 = self.bc
        if gradT_chamber:
            self.bc[0] = {"derivative": gradT_chamber}            
        if T_oven:
            self.bc[1] = {"value": T_oven}

    def evolution_rate(self, state, t=0):
        """implement the python version of the evolution equation"""
        state_lap = state.laplace(bc=self.bc) # , backend="auto")
        # state_grad = state.gradient(bc=self.bc, backend="scipy")
        state_grad2 = state.gradient_squared(bc=self.bc) # , backend="auto")

        '''
        # out_cls_grad2 = state.get_class_by_rank(op_info_grad2.rank_out)
        out_cls_grad = state.get_class_by_rank(op_info_grad.rank_out)
        out_cls_lap = state.get_class_by_rank(op_info_lap.rank_out)
        # state_grad2 = out_cls_grad2(state.grid, data="empty", dtype=state.dtype)
        state_grad = out_cls_grad(state.grid, data="empty", dtype=state.dtype)
        state_lap = out_cls_lap(state.grid, data="empty", dtype=state.dtype)
        state.set_ghost_cells(self.bc)
        # op_grad2(state._data_full, state_grad2.data)
        op_grad(state._data_full, state_grad.data)
        op_lap(state._data_full, state_lap.data)
        '''
        k =  self.kCoef1  * state + self.kCoef0
        cp = self.cpCoef1 * state + self.cpCoef0
        
        state_grad_k_grad = self.kCoef1 * state_grad2 # state_grad.dot(state_grad)
        
        return (state_grad_k_grad + k * state_lap) / cp / self.rho

    '''
    def _make_pde_rhs_numba(self, state):
        """implement the python version of the evolution equation"""
        lap = state.grid.make_operator("laplace", bc=self.bc)
        # grad = state.grid.make_operator("gradient", bc=self.bc)
        grad2 = state.grid.make_operator("gradient_squared", bc=self.bc)
        rho = self.rho
        kCoef0 = self.kCoef0
        kCoef1 = self.kCoef1
        cpCoef0 = self.cpCoef0
        cpCoef1 = self.cpCoef1
        @nb.jit
        def pde_rhs(data, t):
            return (((kCoef1*grad2(data)) + (kCoef1*data + kCoef0)*lap(data)) / rho / (cpCoef1 * data + cpCoef0))
        
        return pde_rhs
'''

class RefractoryWall:
    def __init__ (self, T0):
        self.T_oven = T0
        self.T_chamber = T0 
        self.q_chamber = 0.
        self.T_internal = pde.ScalarField(wall_grid, T0)
        self.eqn = CokeOvenBrickHeatEqn(bc=[{"derivative": 0}, {"value": self.T_oven}])

    def update_bc (self, Q=None, T_oven=None):
        # Q = - k(T) gradT
        # T_chamber = self.T_internal.get_boundary_values(axis=0, upper=False, bc=self.eqn.bc)
        k0 = self.eqn.k(self.T_chamber)
        if Q:
            gradT = Q / wall_area / k0
        else:
            gradT = None
        self.eqn.update_bc(gradT, T_oven)

    def solve (self, dt):
        # solution = self.eqn.solve (eqn, bc)
        self.T_internal = self.eqn.solve(
            self.T_internal, t_range=dt, dt=30., tracker='consistency', backend="numpy")
        self.T_chamber = self.T_internal.get_boundary_values(axis=0, upper=False, bc=self.eqn.bc)

    def heat_to_oven (self):
        """ NOT YET IMPLEMENTED """
        return 0.0


Twall_table = np.loadtxt('./CokeOvenWallTemperature.csv', delimiter=',').T
Twall_table[0] *= ((66/80) / (100/80))
Twall_table[1] += 273.15


def Twall_model(x):
    ''' 
    Coke oven wall temperature vs time after charging 
    Temperature (K) vs Elapsed time (hour)
    '''
    return np.interp(x, Twall_table[0], Twall_table[1])

class OvenChamber:
    def __init__ (self):
        self.content = None

    def get_charge_temperature (self, t):
        """ Return temperature of coal charge content at oven wall """
        if self.content:
            elapsed_time = t - self.content.t_charge
        else:
            elapsed_time = 0.
        return Twall_model(elapsed_time)

    def bake (self, q):
        """ Add transferred heat to coal charge content """
        if self.content:
            self.content.bake(q)

    def charge (self, coal_charge):
        """ Update content with fresh coal is charged """
        self.content = coal_charge

def wall_solve_wrapper(t_range, wall):
    wall.solve(t_range)
    return wall.T_internal, wall.T_chamber

class Battery:
    
    def load_state(self):
        with open('gas.history', 'rb') as gas_history_file:
            self.gas_t_history = pickle.load(gas_history_file)

        with open('wall.history', 'rb') as wall_history_file:
            self.wall_t_history = pickle.load(wall_history_file)

        with open('coke.history', 'rb') as coke_history_file:
            self.product = pickle.load(coke_history_file)

        with open('oven.state', 'rb') as coke_state_file:
            self.processing = pickle.load(coke_state_file)

    def __init__ (self, name, size, heat_program, charge_program, burned_gas_state, hv, init_from_file=False):
        self.name = name                        # Battery name
        self.size = size                        # Size of battery, number of ovens
        self.heat_program = heat_program        # Heat program or schedule object
        self.charge_program = charge_program    # Charge program of schedule object
        self.t = 0                              # Battery time
        self.t_last = 0                         # Time of last Push/Charge
        self.processing = []                    # List of Coke charges under processing(drying)
        self.product = []                       # List of Coke charges done(completed)
        self.gas = ct.Solution('gri30.xml')
        self.gas.TPX = burned_gas_state         # Burned gas T, P, X
        T0, P0, X0 = self.gas.TPX
        self.T0 = T0
        self.P0 = P0
        self.X0 = X0
        self.sequence_idx = 0                   # Integer, 0 ~ (size-1), progress index for oven sequence array

        self.wall_t_history = []
        self.gas_t_history = []

        self.hv = hv # Base unit heat J/kg
        self.normal_heat = self.heat_program.f(-1) # GJ / rev

        # Energy input to battery
        Q0 = self.normal_heat * 1e9 * 3 / 3600 # GJ/rev => J/s (W)
        # Equivalent Fuel+Air mass flow
        mdot0 = Q0 / hv # (J/s) / (J/kg) => kg/s
        self.mdot0 = mdot0 # kg / s

        # chambers[0] - walls_0[0] - ovens[0] - walls_1[0] - chambers[1] - walls_0[1] - ...
        # ... walls_1[i-1] - chambers[i] - walls_0[i] - ovens[i] - walls_1[i] - chambers[i+1] - walls_0[i+1] - ...
        # ... walls_1[size-2] - chambers[size-1] - walls_0[size-1] - ovens[size-1] - walls_1[size-1] - chambers[size]

        self.chambers = [
            CombustionChamber(self.mdot0/self.size, self.gas,
                              (self.T0, self.P0, self.X0), hA=700)
            for ichamber in range(self.size+1)
        ]
        self.ovens = [
            OvenChamber()
            for ioven in range(self.size)
        ]
        self.walls_0 = [
            RefractoryWall(Twall_model(0))
            for ioven in range(self.size)
        ]
        self.walls_1 = [
            RefractoryWall(Twall_model(0))
            for ioven in range(self.size)
        ]
        
        # For 1~4 Coke Ovens with n+5 P/C sequence
        start_indices = [1, 3, 5, 2, 4]
        self.oven_idx_order = np.concatenate([np.array(range(i0 - 1, self.size, 5)) for i0 in start_indices])

        if init_from_file:
            print("Initializaton from file")

            self.load_state()
            latest_chamber = self.gas_t_history[-1]
            latest_wall = self.wall_t_history[-1]

            # Last Record Time
            self.t = latest_chamber[0]
            self.t_last = self.processing[-1].t_charge

            # Recover Chamber State
            for chmbr, T1 in zip(self.chambers, latest_chamber[1]):
                chmbr.T1 = T1

            # Recover Wall State
            for wl, wu, wallT in zip(self.walls_0, self.walls_1, latest_wall[1]):
                wl.T_chamber, wl.T_internal.data, wl.T_oven, wu.T_oven, wu.T_internal.data, wu.T_chamber = wallT

            # Recover Oven State
            for coal in self.processing:
                self.ovens[coal.idx_oven].content = coal
                
        else:
            print("Initializaton Start")
            # 정상 상태 만들기: 모든 문에 n_cycle 회 장입
            n_cycle = 3             # 모든 문 장입 반복 횟수
            period_over_dt = 6.    # period/dt, 장입 간격 / 초기화 time step 크기
            normal_period = self.charge_program.period(-1)  # 감산 전 장입 간격 (주기)
            dt =  normal_period / period_over_dt            # Simulation Time Step
            
            self.t = - normal_period * self.size * n_cycle  # 정상상태 생성 모사 시간 = 장입 간격 * 총 장입 횟수
            self.t_last = self.t                            # 마지막 장입을 정상상태 시뮬레이션 시작 시각으로 설정
        
            # initialization time loop
            for i in range(int(np.ceil(self.size * period_over_dt * n_cycle))):
            # for i in range(3):
                """ Fill battety with normal charge rate """
                self.update(dt)     # Time adavancement

    def mdot (self, t):
        return self.mdot0 * self.heat_program.f(t) / self.normal_heat

    def next_oven (self):
        ''' Index of the oven to which apply push and charge '''
        next_oven_id = self.oven_idx_order[self.sequence_idx % self.size]
        self.sequence_idx += 1
        return next_oven_id

    def bake (self, dt):
        # update combustion chamber equilibrium temperature
            # Tad = 연료 조성과 공연비로 결정
            # m_dot = 연료 발열량과 공급열량 공연비로 결정
            # m(h1 - h0) = hA(Tgas - Twall) => solve with initial T0 = Tad

        # Loop all combustion chambers 
        # update chamber wall temperatures and mass flow rates
        # solve for equilibrium heat to walls
        for i_chamber, chmbr in enumerate(self.chambers):
            if i_chamber > 0:
                wall_lower = self.walls_1[i_chamber-1]
            else:
                wall_lower = None

            if i_chamber < self.size:
                wall_upper = self.walls_0[i_chamber]
            else:
                wall_upper = None

            chmbr.update_mdot(self.mdot(self.t)/self.size)
            chmbr.update_Twall(
                Twall0=(wall_lower.T_chamber if wall_lower else wall_upper.T_chamber), 
                Twall1=(wall_upper.T_chamber if wall_upper else wall_lower.T_chamber), 
            )
            print(f"t={self.t:6.2} : {chmbr.Twall0} K | Chamber {i_chamber} | {chmbr.Twall1} K ")
            chmbr.solve()
            Q1, Q2 = chmbr.heat() # W (J/s)
            if wall_lower: wall_lower.update_bc(Q=Q1)
            if wall_upper: wall_upper.update_bc(Q=Q2)
        
        # Loop all ovens
        # update oven wall temperatures using coke charge age
        # solve heat equations of all walls
        # bake charge in oven
        for i_oven, (oven, wall_lower, wall_upper) in enumerate(zip(self.ovens, self.walls_0, self.walls_1)):
            T_oven = oven.get_charge_temperature(self.t)

            wall_lower.update_bc(T_oven=T_oven)
            wall_upper.update_bc(T_oven=T_oven)

        with Pool(12) as pool:
            wall_sln = pool.starmap(wall_solve_wrapper, [((dt*60*60), w) for w in self.walls_0+self.walls_1])

            # wall_lower.solve(dt * 60 * 60) # convert hours to seconds
            # wall_upper.solve(dt * 60 * 60) # convert hours to seconds
        for ws, wall in zip(wall_sln, self.walls_0+self.walls_1):
            T_internal, T_chamber = ws
            wall.T_internal = T_internal
            wall.T_chamber = T_chamber

        '''
            ql = wall_lower.heat_to_oven()
            qu = wall_upper.heat_to_oven()

            oven.bake(ql+qu) 
            '''

        # advance time oven brick
            # from chamber heat flux boundary condition
            # to oven fixed temperature boundary condition
        
        # integrate heat to oven # 오븐 벽면 온도 우선 시간 함수로

    def push_and_charge (self, coke_charge):
        if len(self.processing) >= self.size:
            self.push(coke_charge.t_charge)
        self.charge(coke_charge)
        
    def push (self, t):
        """ Push complete coke out of oven """
        coke = self.processing.pop(0)
        coke.end_baking(t)
        self.product.append(coke)

    def charge (self, coke_charge):
        self.ovens[coke_charge.idx_oven].charge(coke_charge)
        self.processing.append(coke_charge)

    def dQ (self, dt):
        return self.heat_program.dQ(self.t, self.t+dt)
    
    def is_pc_time (self, dt):
        ''' Whether P/C should be done in this time step '''
        period = self.charge_program.period(self.t)
        return self.t + dt >= period + self.t_last

    def update (self, dt):
        # dQ = self.heat_program.dQ(self.t, self.t+dt) # t, t+dt 사이 공급하는 열량, array 로 대체 필요
        
        # t 에서 t+dt 까지 탄화실 가열
        self.bake(dt)
        
        period = self.charge_program.period(self.t) # 현재 장입 시간 간격
        
        # 마지막 장입탄 장입 시각
        latest_coke_charge = self.processing[-1].t_charge if len(self.processing) > 0 else self.t_last
        
        # t_last + period 가 t, t + dt 사이에 들어오는 것 검사
        # t + dt 가 다음 추출/장입 시각 이후일 때 => 이번 time step 에 추출/장입을 실행해야함
        if self.t + dt >= period + self.t_last :
            # 마지막 장입 시각 + 장입 시간 간격 이 이번 time step 에 포함됨
            # 일정한 간격으로 장입 진행 중, 마지막 장입 시간 += 장입 간격
            if self.t < self.t_last + period:
                self.t_last += period
            # 마지막 장입 이후 현재 장입 간격보다 긴 시간이 경과함 (장입 간격이 짧아짐; 감산 끝남 등)
            # 이번 time step 끝을 마지막 장입 시각으로 업데이트
            else:
                self.t_last = self.t + dt
            
            # 추출/장입 실행
            i_oven = self.next_oven()
            # oven = self.ovens[i_oven]
            fresh_coal = CokeCharge(self.t + dt, i_oven)
            self.push_and_charge(fresh_coal)

            print(f"On {i_oven} P/C within [ {self.t:7.3} , {self.t + dt:7.3} ].", 
                  f"{self.t + dt - latest_coke_charge:7.3} since last P/C. ", 
                  f"period = {self.charge_program.period(self.t):7.3}",)
            
        
        # 시뮬레이션 시간 업데이트
        self.t += dt

        self.gas_t_history.append((self.t, [chmbr.T1 for chmbr in self.chambers]))
        self.wall_t_history.append((self.t, [(wl.T_chamber, wl.T_internal.data, wl.T_oven, wu.T_oven, wu.T_internal.data, wu.T_chamber) for wl, wu in zip(self.walls_0, self.walls_1)]))

def coke_oven_exhaust_stoichiometry (phi=1.0, return_unburned=False):

    # Define the oxidizer composition, here air with 21 mol-% O2 and 79 mol-% N2
    air = "O2:1,N2:3.762"
    coke_oven_fuel = "H2:6.42, O2:0.39, N2:47.28, CH4:1.79, CO:24.25, CO2:19.72, C2H4:0.13, C2H6:0.04"

    mix = ct.Solution('gri30.xml')

    mix.TP = 25+273.15, ct.one_atm
    mix.set_equivalence_ratio(phi=phi, fuel=coke_oven_fuel, oxidizer=air)

    element_X = {ename: mix.elemental_mole_fraction(ename) for ename in mix.element_names}
    
    exhaust = ct.Solution('gri30.xml')
    exhaust.TPX = (25+273.15, ct.one_atm,
                   {
                       "CO2": element_X['C'],
                       "H2O": element_X['H']/2,
                       "O2": (element_X['O'] - 2*element_X['C'] - element_X['H']/2)/2,
                       "N2": element_X['N']/2,
                   }
                   )

    if return_unburned:
        return mix.mole_fraction_dict(threshold=-1), exhaust.mole_fraction_dict(threshold=-1)
    else:
        return exhaust.mole_fraction_dict(threshold=-1)

class HeatSchedule:
    def __init__ (self, xp, fp):
        self.xp = xp
        self.fp = fp
        self.f = lambda x: np.interp(x, self.xp, self.fp)
        
    def dQ(self, t0, t1):
        x = np.linspace(t0, t1, 31)
        return np.trapz(self.f(x), x)        

class ChargeSchedule:
    def __init__ (self, normal_load, service_start, service_time, service_load, aux_start, aux_time, aux_load):
        self.xp = np.array([service_start, service_start, service_start+service_time, service_start+service_time, 
                            aux_start, aux_start, aux_start+aux_time, aux_start+aux_time, ])
        self.fp = np.array([normal_load, service_load, service_load, normal_load,
                            normal_load, aux_load, aux_load, normal_load])
        self.f = lambda x: np.interp(x, self.xp, self.fp)
        
    def to_charge (self, t0, t1):
        self.f(t0)
        return np.trapz(self.f(x), x)
    
    def period (self, t):
        return 24 / self.f(t)

if __name__ == "__main__":

    # Define the oxidizer composition, here air with 21 mol-% O2 and 79 mol-% N2
    air = "O2:0.21,N2:0.79"
    coke_oven_fuel = "H2:6.42, O2:0.39, N2:47.28, CH4:1.79, CO:24.25, CO2:19.72, C2H4:0.13, C2H6:0.04"

    f_found = optimize.root_scalar(lambda x: coke_oven_exhaust_stoichiometry(x)["O2"] - 0.045, 
                          bracket=[1e-300, 1])

    # equivalence ratio for O2 4.5 % in exhaust gas (stoichiometric)
    phi_O2_045 = f_found.root

    # unburned and burned gas compositions for O2 4.5 % in exhaust gas (stoichiometric)
    Xu, Xb = coke_oven_exhaust_stoichiometry(phi_O2_045, return_unburned=True)

    gas = ct.Solution('gri30.xml')

    # Heating value of unburned premixed gas
    gas.TPX = 25 + 273.15, ct.one_atm, Xu
    hu = gas.enthalpy_mass
    gas.TPX = None, None, Xb
    hb = gas.enthalpy_mass
    hv = hu - hb
    print(f'{hu*1e-6} - {hb*1e-6} = {hv*1e-6} MJ/kg')

    # burned premixed gas state (chemical equilibrium with HP constraint)
    gas.TP = 600+273.15, ct.one_atm
    gas.set_equivalence_ratio(phi_O2_045, fuel=coke_oven_fuel, oxidizer=air)
    gas.equilibrate('HP')
    gas_in_state = gas.TPX

    # Time(Hours) - GJ/rev
    sample_program = np.array(open('sample_heat_221128_3A.txt').read().split(), dtype=np.double).reshape((-1,2))

    heating_plan = HeatSchedule(*sample_program.T)
    charging_plan = ChargeSchedule( 82, 9, 13, 1e-12, 9+13+18, 4, 1e-12 )
    n_doors = 66

    load_state = False
    bat3A = Battery("3A", n_doors, heating_plan, charging_plan, gas_in_state, hv, init_from_file=load_state)

    if not load_state:
        with open('gas.history', 'wb') as gas_history_file:
            pickle.dump(bat3A.gas_t_history, gas_history_file)

        with open('wall.history', 'wb') as wall_history_file:
            pickle.dump(bat3A.wall_t_history, wall_history_file)

        with open('coke.history', 'wb') as wall_history_file:
            pickle.dump(bat3A.product, wall_history_file)

        with open('oven.state', 'wb') as wall_history_file:
            pickle.dump(bat3A.processing, wall_history_file)

    dt = 5. * 1./60. # 5 min
    for it in range (int(60/dt)): # 시뮬레이션 시간 도메인 = 60시간
        bat3A.update(dt)

    with open('gas.history2', 'wb') as gas_history_file:
        pickle.dump(bat3A.gas_t_history, gas_history_file)

    with open('wall.history2', 'wb') as wall_history_file:
        pickle.dump(bat3A.wall_t_history, wall_history_file)

    with open('coke.history2', 'wb') as wall_history_file:
        pickle.dump(bat3A.product, wall_history_file)

    with open('oven.state2', 'wb') as wall_history_file:
        pickle.dump(bat3A.processing, wall_history_file)

    print("Done")