"""Coke Oven Battery Simulation Module.

This module models the thermal behavior, heat transfer, and scheduling operations of
a coke oven battery. It implements systems of flues (combustion chambers), refractory
brick walls (solving 1D heat equations), and oven chambers with loaded coal charges.
"""

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 cantera as ct
from scipy import optimize

# PDE Solver Configuration
USE_CUSTOM_SOLVER = True  # Set to False to use the baseline py-pde solver for verification

try:
    import pde
except ImportError:
    pde = None


class CombustionChamber:
    """Represents a combustion chamber in the coke oven battery.

    This class models the steady-state thermal energy balance of combustion gases
    flowing through the heating flues (chambers) adjacent to the ovens.

    Attributes:
        mdot (float): Mass flow rate of fuel-air mixture (kg/s).
        gas (cantera.Solution): Cantera gas object.
        eq_state (tuple): Thermodynamic state (T, P, X) of burned gas.
        T0 (float): Adiabatic flame temperature (K).
        P0 (float): Operating pressure (Pa).
        X0 (dict/array): Fuel-air mole fractions.
        h0 (float): Inlet enthalpy of the gas (J/kg).
        hA (float): Heat transfer coefficient times area (W/K).
        T1 (float): Outlet gas temperature (K).
        Twall0 (float): Temperature of the lower wall (K).
        Twall1 (float): Temperature of the upper wall (K).
        Area (float): Oven cross section area (m^2).
    """

    def __init__(self, mdot, ct_object, burned_state, hA=700):
        """Initializes the CombustionChamber.

        Args:
            mdot (float): Mass flow rate (kg/s).
            ct_object (cantera.Solution): Instantiated Cantera Solution object.
            burned_state (tuple): State variables (T, P, X) representing burned gas.
            hA (float, optional): Heat transfer coefficient * area. Defaults to 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):
        """Updates the mass flow rate if a new value is provided.

        Args:
            mdot_new (float): The new mass flow rate (kg/s).
        """
        if mdot_new:
            self.mdot = mdot_new

    def update_Twall(self, Twall0=None, Twall1=None):
        """Updates the boundary wall temperatures.

        Args:
            Twall0 (float, optional): Lower wall temperature (K).
            Twall1 (float, optional): Upper wall temperature (K).
        """
        if Twall0:
            self.Twall0 = Twall0
        if Twall1:
            self.Twall1 = Twall1

    def energy_balance_equation(self, Tout):
        """Calculates the residual energy imbalance for root-finding.

        Args:
            Tout (float): Guessed outlet gas temperature (K).

        Returns:
            float: Energy balance residual (W).
        """
        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 solves for the outlet temperature balancing heat loss to walls.

        Returns:
            float: Resolved outlet temperature (K).
        """
        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 meanTwall

        return f_found.root

    def heat(self, Tout=None):
        """Calculates the heat transfer rate to the walls.

        Args:
            Tout (float, optional): Outlet gas temperature (K). If None, uses T1.

        Returns:
            tuple: Heat transfer rates (q0, q1) to the lower and upper walls (W).
        """
        if Tout is None:
            Tout = self.T1
        Tgas = (self.T0 + Tout) / 2

        return self.hA * (Tgas - self.Twall0), self.hA * (Tgas - self.Twall1)


class Regenerator:
    """Models a single regenerator chamber in the coke oven battery.

    Uses a 1D lumped thermal mass model to track temperature dynamics during
    heat storage (charging) and release (discharging) cycles.
    """

    def __init__(self, idx, T_init=1016.9, M_reg=1e9, Cp_reg=1000.0, eff=0.8):
        """Initializes the Regenerator.

        Args:
            idx (int): Index identifier of the regenerator.
            T_init (float, optional): Initial brick temperature (K). Default is 1016.9 K
                to replicate the 600°C boundary condition.
            M_reg (float, optional): Mass of checker bricks (kg). Default is 1e9
                (option A: infinite thermal mass to suppress dynamics for verification).
            Cp_reg (float, optional): Specific heat of checker bricks (J/kg/K). Default 1000.
            eff (float, optional): Heat exchange effectiveness (0 to 1). Default 0.8.
        """
        self.idx = idx
        self.T = T_init
        self.M_reg = M_reg
        self.Cp_reg = Cp_reg
        self.eff = eff
        self.Cp_gas = 1000.0  # J/kg/K, average heat capacity of gas mixture

    def heat_exchange(self, dt_sec, mdot, T_gas_in, is_discharging):
        """Calculates gas outlet temperature and updates regenerator brick temperature.

        Args:
            dt_sec (float): Time step in seconds.
            mdot (float): Mass flow rate of air/flue gas (kg/s).
            T_gas_in (float): Temperature of incoming gas (K).
            is_discharging (bool): True if air/gas is being preheated (discharging),
                False if hot flue gas is reheating the bricks (charging).

        Returns:
            float: Temperature of the outgoing gas (K).
        """
        if mdot <= 1e-8:
            return T_gas_in

        if is_discharging:
            # Heat release cycle: incoming cold air at T_gas_in is preheated to T_gas_out
            T_gas_out = T_gas_in + self.eff * (self.T - T_gas_in)
            Q_dot = mdot * self.Cp_gas * (T_gas_out - T_gas_in)  # W (J/s)
            self.T -= (Q_dot * dt_sec) / (self.M_reg * self.Cp_reg)
            return T_gas_out
        else:
            # Heat storage cycle: incoming hot flue gas at T_gas_in heats the bricks
            T_gas_out = T_gas_in - self.eff * (T_gas_in - self.T)
            Q_dot = mdot * self.Cp_gas * (T_gas_in - T_gas_out)  # W (J/s)
            self.T += (Q_dot * dt_sec) / (self.M_reg * self.Cp_reg)
            return T_gas_out


class CokeCharge:
    """Represents a single coal/coke charge inside an oven.

    Attributes:
        t_charge (float): Simulation timestamp when coal was charged (hours).
        t_push (float or None): Simulation timestamp when coke was pushed (hours).
        idx_oven (int): Index of the oven this charge belongs to.
        Q (float): Total heat absorbed by this charge (J).
    """

    def __init__(self, t_charge, idx_oven):
        """Initializes CokeCharge.

        Args:
            t_charge (float): Time of charging.
            idx_oven (int): Target oven index.
        """
        self.t_charge = t_charge
        self.t_push = None
        self.idx_oven = idx_oven
        self.Q = 0

    def bake(self, dQ):
        """Applies a increment of thermal energy to the charge.

        Args:
            dQ (float): Energy increment (J).
        """
        self.Q += dQ

    def end_baking(self, t):
        """Finalizes the charge lifecycle at push time.

        Args:
            t (float): Pushing time (hours).
        """
        self.t_push = t


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


class TInternal:
    """Helper wrapper around numerical internal temperature data for a wall.

    Attributes:
        data (numpy.ndarray): Spatial temperature profile within the brick wall.
    """

    def __init__(self, data):
        """Initializes TInternal wrapper.

        Args:
            data (array-like): Numerical temperature array.
        """
        self.data = np.array(data, dtype=np.float64)

    def get_boundary_values(self, axis=0, upper=False, bc=None):
        """Extracts boundary temperature values based on boundary conditions.

        Args:
            axis (int, optional): Spatial axis (default is 0).
            upper (bool, optional): If True, gets right-hand boundary (oven-side),
                otherwise left-hand boundary (chamber-side). Defaults to False.
            bc (list, optional): Boundary conditions list.

        Returns:
            float: The computed boundary temperature value (K).
        """
        dx = brick_thickness / n_grid_brick
        if not upper:
            g_L = 0.0
            if bc and len(bc) > 0:
                bc0 = bc[0]
                if isinstance(bc0, dict):
                    g_L = bc0.get("derivative", 0.0)
                else:
                    g_L = bc0
            return self.data[0] + 0.5 * dx * g_L
        else:
            T_R = 0.0
            if bc and len(bc) > 1:
                bc1 = bc[1]
                if isinstance(bc1, dict):
                    T_R = bc1.get("value", 0.0)
                else:
                    T_R = bc1
            return T_R


class CokeOvenBrickHeatEqnBase:
    """Base class defining physical parameters for the brick wall heat equation.

    Attributes:
        bc (list): Boundary conditions.
        rho (float): Density of refractory brick (kg/m3).
        kCoef0 (float): Constant coefficient of thermal conductivity (W/m/K).
        kCoef1 (float): Temperature coefficient of thermal conductivity (W/m/K2).
        cpCoef0 (float): Constant coefficient of specific heat (J/kg/K).
        cpCoef1 (float): Temperature coefficient of specific heat (J/kg/K2).
    """

    def __init__(self, bc="auto_periodic_neumann"):
        """Initializes CokeOvenBrickHeatEqnBase.

        Args:
            bc (str or list, optional): Boundary conditions description.
                Defaults to "auto_periodic_neumann".
        """
        try:
            super().__init__()
        except Exception:
            pass
        self._cache = {}
        if bc == "auto_periodic_neumann":
            self.bc = [{"derivative": 0.0}, {"value": 0.0}]
        else:
            self.bc = list(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):
        """Calculates temperature-dependent thermal conductivity.

        Args:
            T (float or numpy.ndarray): Temperature (K).

        Returns:
            float or numpy.ndarray: Thermal conductivity (W/m/K).
        """
        return T * self.kCoef1 + self.kCoef0

    def cp(self, T):
        """Calculates temperature-dependent specific heat capacity.

        Args:
            T (float or numpy.ndarray): Temperature (K).

        Returns:
            float or numpy.ndarray: Specific heat capacity (J/kg/K).
        """
        return T * self.cpCoef1 + self.cpCoef0

    def update_bc(self, gradT_chamber=None, T_oven=None):
        """Updates boundary condition parameters.

        Args:
            gradT_chamber (float, optional): Temperature gradient at chamber side.
            T_oven (float, optional): Oven side boundary temperature (K).
        """
        if gradT_chamber is not None:
            self.bc[0] = {"derivative": gradT_chamber}
        if T_oven is not None:
            self.bc[1] = {"value": T_oven}


if pde is not None:
    class CokeOvenBrickHeatEqn(CokeOvenBrickHeatEqnBase, pde.PDEBase):
        """1D Heat Equation model for coke oven brick walls, leveraging py-pde package."""

        def evolution_rate(self, state, t=0):
            """Calculates time derivative of temperature field for solvers.

            Args:
                state (pde.ScalarField): Current temperature field.
                t (float): Current simulation time.

            Returns:
                pde.ScalarField: Evolution rate dT/dt.
            """
            state_lap = state.laplace(bc=self.bc)
            state_grad2 = state.gradient_squared(bc=self.bc)
            k = self.kCoef1 * state + self.kCoef0
            cp = self.cpCoef1 * state + self.cpCoef0
            state_grad_k_grad = self.kCoef1 * state_grad2
            return (state_grad_k_grad + k * state_lap) / cp / self.rho
else:
    class CokeOvenBrickHeatEqn(CokeOvenBrickHeatEqnBase):
        """Fallback heat equation model without py-pde dependencies."""
        pass


class RefractoryWall:
    """Simulates a refractory brick wall separating combustion chambers and ovens.

    Solves the 1D transient heat conduction equation through the refractory wall
    using either py-pde or a custom NumPy finite difference solver.

    Attributes:
        T_oven (float): Temperature at the oven side (K).
        T_chamber (float): Temperature at the combustion chamber side (K).
        q_chamber (float): Heat flux from chamber.
        T_internal (TInternal or pde.ScalarField): Internal temperature field.
        eqn (CokeOvenBrickHeatEqn): Heat equation PDE model instance.
    """

    def __init__(self, T0):
        """Initializes RefractoryWall.

        Args:
            T0 (float): Initial uniform temperature (K).
        """
        self.T_oven = T0
        self.T_chamber = T0
        self.q_chamber = 0.
        if USE_CUSTOM_SOLVER:
            self.T_internal = TInternal(np.full(n_grid_brick, T0))
        else:
            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):
        """Updates the wall boundary conditions based on energy flow.

        Args:
            Q (float, optional): Heat input rate (W).
            T_oven (float, optional): Oven boundary temperature (K).
        """
        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):
        """Solves the heat equation over time interval dt.

        Args:
            dt (float): Simulation time step (seconds).
        """
        if USE_CUSTOM_SOLVER:
            dt_internal = 30.0
            steps = int(round(dt / dt_internal))
            dx = brick_thickness / n_grid_brick
            
            T = self.T_internal.data
            
            g_L = 0.0
            if self.eqn.bc and len(self.eqn.bc) > 0:
                bc0 = self.eqn.bc[0]
                if isinstance(bc0, dict):
                    g_L = bc0.get("derivative", 0.0)
                else:
                    g_L = bc0
                    
            T_R = 0.0
            if self.eqn.bc and len(self.eqn.bc) > 1:
                bc1 = self.eqn.bc[1]
                if isinstance(bc1, dict):
                    T_R = bc1.get("value", 0.0)
                else:
                    T_R = bc1
            
            for _ in range(steps):
                T_minus_1 = T[0] + dx * g_L
                T_N = 2.0 * T_R - T[-1]
                
                T_aug = np.empty(n_grid_brick + 2)
                T_aug[0] = T_minus_1
                T_aug[1:-1] = T
                T_aug[-1] = T_N
                
                grad = (T_aug[2:] - T_aug[:-2]) / (2.0 * dx)
                grad2 = grad * grad
                lap = (T_aug[2:] - 2.0 * T_aug[1:-1] + T_aug[:-2]) / (dx * dx)
                
                k = self.eqn.kCoef1 * T + self.eqn.kCoef0
                cp = self.eqn.cpCoef1 * T + self.eqn.cpCoef0
                
                dTdt = (self.eqn.kCoef1 * grad2 + k * lap) / (cp * self.eqn.rho)
                T += dt_internal * dTdt
                
            self.T_chamber = T[0] + 0.5 * dx * g_L
        else:
            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):
        """Calculates heat transfer to the oven chamber.

        Returns:
            float: Heat transfer (W). Not implemented yet.
        """
        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):
    """Calculates the coke oven wall temperature based on elapsed time since charging.

    Args:
        x (float): Elapsed time (hours).

    Returns:
        float: Oven wall temperature (K).
    """
    return np.interp(x, Twall_table[0], Twall_table[1])


class OvenChamber:
    """Represents an individual oven chamber containing a coke charge.

    Attributes:
        content (CokeCharge or None): The coke charge model inside the oven.
    """

    def __init__(self):
        """Initializes OvenChamber."""
        self.content = None

    def get_charge_temperature(self, t):
        """Gets the temperature of the coal charge content at the oven wall.

        Args:
            t (float): Simulation time (hours).

        Returns:
            float: Charge surface temperature (K).
        """
        if self.content:
            elapsed_time = t - self.content.t_charge
        else:
            elapsed_time = 0.
        return Twall_model(elapsed_time)

    def bake(self, q):
        """Applies baking heat to the coal charge.

        Args:
            q (float): Heat energy applied (J).
        """
        if self.content:
            self.content.bake(q)

    def charge(self, coal_charge):
        """Charges fresh coal into the oven chamber.

        Args:
            coal_charge (CokeCharge): The coal charge object to load.
        """
        self.content = coal_charge


def wall_solve_wrapper(t_range, wall):
    """Worker function wrapper to solve wall heat equation in parallel.

    Args:
        t_range (float): Time range to solve (seconds).
        wall (RefractoryWall): Wall object instance to solve.

    Returns:
        tuple: (updated T_internal field, updated boundary T_chamber temperature)
    """
    wall.solve(t_range)
    return wall.T_internal, wall.T_chamber


class Battery:
    """Represents a complete Coke Oven Battery.

    A battery consists of a series of alternating combustion chambers, refractory
    brick walls, and oven chambers, along with corresponding schedules for charging
    and heating.

    Attributes:
        name (str): Battery name identifier.
        size (int): Number of oven chambers.
        heat_program (HeatSchedule): Operational heating program schedule.
        charge_program (ChargeSchedule): Operational coal charging schedule.
        t (float): Current simulation time (hours).
        t_last (float): Timestamp of the last Push/Charge event (hours).
        processing (list of CokeCharge): Currently active coke charges.
        product (list of CokeCharge): Log of completed coke charges.
        gas (cantera.Solution): Local Cantera Solution object.
        T0 (float): Adiabatic flame temperature of incoming gas (K).
        P0 (float): Gas operating pressure (Pa).
        X0 (dict): Gas composition.
        sequence_idx (int): Current sequence progress index.
        wall_t_history (list): Recorded history of wall temperatures.
        gas_t_history (list): Recorded history of chamber temperatures.
        hv (float): Heating value of fuel-air mix (J/kg).
        normal_heat (float): Baseline heat load (GJ/rev).
        mdot0 (float): Baseline fuel mixture mass flow rate (kg/s).
        chambers (list of CombustionChamber): Combustion flues.
        ovens (list of OvenChamber): Oven chambers.
        walls_0 (list of RefractoryWall): Lower refractory walls.
        walls_1 (list of RefractoryWall): Upper refractory walls.
        oven_idx_order (numpy.ndarray): Charging schedule oven sequence.
    """

    def load_state(self):
        """Loads simulation state from binary history files."""
        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):
        """Initializes Battery simulation.

        Args:
            name (str): Identifier name.
            size (int): Total count of ovens.
            heat_program (HeatSchedule): Heating scheduler object.
            charge_program (ChargeSchedule): Charging scheduler object.
            burned_gas_state (tuple): Initial TPX state of burned flue gas.
            hv (float): Net heating value (J/kg).
            init_from_file (bool, optional): Recover state from pickle. Defaults to 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
        # List of Coke charges under processing(drying)
        self.processing = []
        # List of Coke charges done(completed)
        self.product = []
        self.gas = ct.Solution('gri30.yaml')
        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
        # Integer, 0 ~ (size-1), progress index for oven sequence array
        self.sequence_idx = 0

        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)
        ]

        # Regenerators and fuel valves initialization (size + 2 elements)
        # Option A is default: M_reg=1e9 to replicate the legacy 600°C boundary condition.
        self.regenerators = [
            Regenerator(ireg, T_init=1016.9, M_reg=1e9, Cp_reg=1000.0, eff=0.8)
            for ireg in range(self.size + 2)
        ]
        self.fuel_valves = np.zeros(self.size + 2)
        self.reversing_period = 20.0 / 60.0  # 20 minutes in hours
        self.control_active = False

        # Cantera parameters for dynamic preheating and equilibrium calculations
        self.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"
        self.oxidizer = "O2:0.21,N2:0.79"
        try:
            f_found = optimize.root_scalar(
                lambda x: coke_oven_exhaust_stoichiometry(x)["O2"] - 0.045,
                bracket=[1e-300, 1]
            )
            self.phi = f_found.root
        except Exception:
            self.phi = 0.814238515

        # 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 with copy() to break reference linkage with history lists
            for wl, wu, wallT in zip(self.walls_0, self.walls_1, latest_wall[1]):
                wl.T_chamber = wallT[0]
                wl.T_internal.data = wallT[1].copy()
                wl.T_oven = wallT[2]
                wu.T_oven = wallT[3]
                wu.T_internal.data = wallT[4].copy()
                wu.T_chamber = wallT[5]

            # 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):
        """Calculates mass flow rate of gas at time t.

        Args:
            t (float): Simulation time (hours).

        Returns:
            float: Mass flow rate (kg/s).
        """
        return self.mdot0 * self.heat_program.f(t) / self.normal_heat

    def is_cycle_A(self, t):
        """Checks if the battery is currently in Cycle A (odd-numbered regenerators discharging).

        Args:
            t (float): Simulation time (hours).

        Returns:
            bool: True if in Cycle A, False if in Cycle B.
        """
        cycle_num = int(np.floor(t / self.reversing_period))
        return (cycle_num % 2) == 0

    def get_chamber_inlets(self, t):
        """Computes mass flow rate and maps regenerators for each combustion chamber.

        Args:
            t (float): Current simulation time (hours).

        Returns:
            tuple: (chamber_mdots, inlet_reg_indices, outlet_reg_indices)
                - chamber_mdots (np.ndarray): Fuel-air flow rate for each chamber (kg/s).
                - inlet_reg_indices (list): Regenerator index supplying to each chamber.
                - outlet_reg_indices (list): Regenerator index receiving exhaust from each chamber.
        """
        total_mdot = self.mdot(t)
        size = self.size

        # If external control is inactive, dynamically reset valves to legacy distribution
        if not getattr(self, "control_active", False):
            self.fuel_valves = np.zeros(size + 2)
            if self.is_cycle_A(t):
                # Cycle A: Odd regenerators (0-indexed even j) discharge
                self.fuel_valves[0] = total_mdot / size
                self.fuel_valves[2:size+1:2] = 2.0 * total_mdot / size
            else:
                # Cycle B: Even regenerators (0-indexed odd j) discharge
                self.fuel_valves[1:size+1:2] = 2.0 * total_mdot / size
                self.fuel_valves[size+1] = total_mdot / size

        # Distribute valve flows to individual combustion chambers
        chamber_mdots = np.zeros(size + 1)
        inlet_reg_indices = [-1] * (size + 1)
        outlet_reg_indices = [-1] * (size + 1)

        is_a = self.is_cycle_A(t)

        if is_a:
            # Cycle A: Even regenerators (j = 0, 2, ..., size) are INLETS
            # Odd regenerators (j = 1, 3, ..., size+1) are OUTLETS
            for j in range(0, size + 1, 2):
                val = self.fuel_valves[j]
                if j == 0:
                    chamber_mdots[0] += val
                    inlet_reg_indices[0] = 0
                else:
                    chamber_mdots[j-1] += val / 2.0
                    chamber_mdots[j] += val / 2.0
                    inlet_reg_indices[j-1] = j
                    inlet_reg_indices[j] = j

            for j in range(1, size + 2, 2):
                if j == size + 1:
                    outlet_reg_indices[size] = size + 1
                else:
                    outlet_reg_indices[j-1] = j
                    outlet_reg_indices[j] = j
        else:
            # Cycle B: Odd regenerators (j = 1, 3, ..., size+1) are INLETS
            # Even regenerators (j = 0, 2, ..., size) are OUTLETS
            for j in range(1, size + 2, 2):
                val = self.fuel_valves[j]
                if j == size + 1:
                    chamber_mdots[size] += val
                    inlet_reg_indices[size] = size + 1
                else:
                    chamber_mdots[j-1] += val / 2.0
                    chamber_mdots[j] += val / 2.0
                    inlet_reg_indices[j-1] = j
                    inlet_reg_indices[j] = j

            for j in range(0, size + 1, 2):
                if j == 0:
                    outlet_reg_indices[0] = 0
                else:
                    outlet_reg_indices[j-1] = j
                    outlet_reg_indices[j] = j

        return chamber_mdots, inlet_reg_indices, outlet_reg_indices

    def next_oven(self):
        """Returns the index of the next oven to be pushed and charged.

        Returns:
            int: Oven index (0-indexed).
        """
        next_oven_id = self.oven_idx_order[self.sequence_idx % self.size]
        self.sequence_idx += 1
        return next_oven_id

    def bake(self, dt):
        """Advances thermal states of combustion chambers, walls, and ovens.

        Args:
            dt (float): Simulation time step (hours).
        """
        dt_sec = dt * 3600.0
        size = self.size
        is_a = self.is_cycle_A(self.t)

        # Get dynamic mass flows and regenerator mappings for this step
        chamber_mdots, inlet_reg_indices, outlet_reg_indices = self.get_chamber_inlets(self.t)

        # 1. Aggregate flow rates for the discharging (inlet) regenerators
        inlet_mdots = np.zeros(size + 2)
        for i_chamber, mdot in enumerate(chamber_mdots):
            inlet_reg_idx = inlet_reg_indices[i_chamber]
            inlet_mdots[inlet_reg_idx] += mdot

        # 2. Perform heat exchange for discharging regenerators once to find preheat temperatures
        reg_preheat_temps = np.zeros(size + 2)
        active_inlets = range(0, size + 1, 2) if is_a else range(1, size + 2, 2)
        for j in active_inlets:
            reg_preheat_temps[j] = self.regenerators[j].heat_exchange(
                dt_sec, inlet_mdots[j], 298.15, is_discharging=True
            )

        # Cache for Cantera HP equilibrium calculation to save CPU overhead.
        # Maps inlet_reg_idx -> (T0, burned_state, h0)
        cantera_cache = {}

        # Arrays to aggregate exhaust properties for the charging (outlet) regenerators
        outlet_mdots = np.zeros(size + 2)
        outlet_T_weighted = np.zeros(size + 2)

        # 3. Loop all combustion chambers to solve thermal equations
        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

            mdot = chamber_mdots[i_chamber]
            inlet_reg_idx = inlet_reg_indices[i_chamber]
            T_pre = reg_preheat_temps[inlet_reg_idx]

            # Compute new combustion state (HP equilibrium) based on T_pre
            if inlet_reg_idx not in cantera_cache:
                self.gas.TP = T_pre, self.P0
                self.gas.set_equivalence_ratio(self.phi, self.fuel, self.oxidizer)
                self.gas.equilibrate('HP')
                burned_state = self.gas.TPX
                h0 = self.gas.enthalpy_mass
                T0 = self.gas.T
                cantera_cache[inlet_reg_idx] = (T0, burned_state, h0)

            T0, burned_state, h0 = cantera_cache[inlet_reg_idx]

            # Update chamber state variables
            chmbr.update_mdot(mdot)
            chmbr.gas.TPX = burned_state
            chmbr.T0 = T0
            chmbr.h0 = h0
            chmbr.T1 = T0  # reset outlet guess to Tad

            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)

            # Accumulate exhaust properties for charging (outlet) regenerator
            outlet_reg_idx = outlet_reg_indices[i_chamber]
            outlet_mdots[outlet_reg_idx] += mdot
            outlet_T_weighted[outlet_reg_idx] += mdot * chmbr.T1

        # 4. Perform heat exchange for charging regenerators once using mass-weighted exhaust temps
        active_outlets = range(1, size + 2, 2) if is_a else range(0, size + 1, 2)
        for j in active_outlets:
            mdot_tot = outlet_mdots[j]
            if mdot_tot > 1e-8:
                T_exh_avg = outlet_T_weighted[j] / mdot_tot
                self.regenerators[j].heat_exchange(
                    dt_sec, mdot_tot, T_exh_avg, is_discharging=False
                )

        # 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)

        if USE_CUSTOM_SOLVER:
            for w in self.walls_0 + self.walls_1:
                w.solve(dt * 60 * 60)
        else:
            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])

            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):
        """Orchestrates pushing older coke and charging fresh coal.

        Args:
            coke_charge (CokeCharge): The fresh coke charge instance.
        """
        if len(self.processing) >= self.size:
            self.push(coke_charge.t_charge)
        self.charge(coke_charge)

    def push(self, t):
        """Pushes the finished coke out of the oven.

        Args:
            t (float): Current time (hours).
        """
        coke = self.processing.pop(0)
        coke.end_baking(t)
        self.product.append(coke)

    def charge(self, coke_charge):
        """Charges a fresh coal unit into the oven list.

        Args:
            coke_charge (CokeCharge): The coal charge instance.
        """
        self.ovens[coke_charge.idx_oven].charge(coke_charge)
        self.processing.append(coke_charge)

    def dQ(self, dt):
        """Calculates total heat supplied over time interval dt.

        Args:
            dt (float): Time interval (hours).

        Returns:
            float: Cumulative heat (GJ).
        """
        return self.heat_program.dQ(self.t, self.t+dt)

    def is_pc_time(self, dt):
        """Checks if push/charge should happen in the current time step.

        Args:
            dt (float): Time step (hours).

        Returns:
            bool: True if push/charge is scheduled.
        """
        period = self.charge_program.period(self.t)
        return self.t + dt >= period + self.t_last

    def update(self, dt):
        """Advances simulation by dt.

        Args:
            dt (float): Simulation step size (hours).
        """
        # 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.copy(), wl.T_oven, wu.T_oven,
                                   wu.T_internal.data.copy(), 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):
    """Calculates exhaust gas composition for coke oven gas combustion.

    Args:
        phi (float, optional): Equivalence ratio. Defaults to 1.0.
        return_unburned (bool, optional): If True, returns both unburned and
            burned gas compositions. Defaults to False.

    Returns:
        dict or tuple: Burned composition dictionary, or (unburned, burned) tuple.
    """
    # 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.yaml')

    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.yaml')
    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:
    """Represents a heat supply program schedule.

    Attributes:
        xp (array-like): Timeline anchor points (hours).
        fp (array-like): Heat loads at anchor points (GJ/rev).
        f (callable): Interpolation function mapping time -> heat load.
    """

    def __init__(self, xp, fp):
        """Initializes HeatSchedule.

        Args:
            xp (array-like): Timeline hours.
            fp (array-like): Heat load array.
        """
        self.xp = xp
        self.fp = fp
        self.f = lambda x: np.interp(x, self.xp, self.fp)

    def dQ(self, t0, t1):
        """Integrates heat input from time t0 to t1.

        Args:
            t0 (float): Start time (hours).
            t1 (float): End time (hours).

        Returns:
            float: Cumulative heat (GJ).
        """
        x = np.linspace(t0, t1, 31)
        return np.trapz(self.f(x), x)


class ChargeSchedule:
    """Represents the scheduling sequence of coal charging operations.

    Attributes:
        xp (numpy.ndarray): Charging program phase change hours.
        fp (numpy.ndarray): Charging rates during phases.
        f (callable): Interpolation function mapping time -> charging rate.
    """

    def __init__(self, normal_load, service_start, service_time, service_load, aux_start, aux_time, aux_load):
        """Initializes ChargeSchedule.

        Args:
            normal_load (float): Baseline charging rate.
            service_start (float): Start hour for maintenance service.
            service_time (float): Duration of maintenance service (hours).
            service_load (float): Charging rate during maintenance.
            aux_start (float): Start hour for auxiliary service phase.
            aux_time (float): Duration of auxiliary phase (hours).
            aux_load (float): Charging rate during auxiliary phase.
        """
        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):
        """Calculates cumulative coal units charged between t0 and t1.

        Args:
            t0 (float): Start time.
            t1 (float): End time.

        Returns:
            float: Total units charged.
        """
        # (Note: 'x' is not defined here in original, keeping it as is to preserve original logic)
        return np.trapz(self.f(x), x)

    def period(self, t):
        """Calculates the time interval between subsequent charges.

        Args:
            t (float): Current time (hours).

        Returns:
            float: Time period (hours).
        """
        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.yaml')

    # 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-Plan2.txt').read().split(), dtype=np.double).reshape((-1, 2))

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

    load_state = True
    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")