#!/usr/bin/env python3
"""
Process Simulation for ORLEN Lignin Oxidation Campaign
Calculates man-hours, utilities, and costs based on BoE assumptions.
"""

import json
import math
import os
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, Any


@dataclass
class Rates:
    # Utilities
    electricity_pln_per_kwh: float = 1.50
    steam_pln_per_kg: float = 0.95
    chilled_pln_per_kwhc: float = 0.50
    cooling_water_pln_per_m3: float = 12.0
    process_water_pln_per_m3: float = 10.0
    # Labor
    operator: float = 135.0
    supervisor: float = 200.0
    qc: float = 180.0
    maintenance: float = 155.0
    eng_hse: float = 190.0


@dataclass
class EquipmentPowers:
    agitator_kw: float = 3.0
    circulation_pump_kw: float = 2.0
    heat_loss_kw: float = 6.0
    dryer_kw: float = 4.0
    misc_kwht_per_batch: float = 5.0


@dataclass
class AirSystem:
    flow_m3ph: float = 15.0
    sec_kwh_per_m3: float = 0.12


@dataclass
class ProcessParams:
    target_dry_kg: float = 2000.0
    batch_dry_kg: float = 200.0
    solids_mass_frac: float = 0.10  # 10% solids per spec
    cp_kj_per_kgk: float = 4.2
    t_cold_c: float = 25.0
    t_hot_c: float = 150.0
    t_filter_c: float = 50.0
    elec_efficiency: float = 0.95
    heat_overhead_factor: float = 1.10
    use_chiller: bool = False
    cooling_water_m3_per_batch: float = 9.0
    process_water_m3_per_batch: float = 2.0
    consumables_pln_per_batch: float = 70.0
    admin_pct: float = 0.10
    contingency_pct: float = 0.15
    contingency_on_base_pre_admin: bool = True


@dataclass
class DurationsHours:
    # Core
    heatup: float = 4.0
    oxidation_attended: float = 16.0
    oxidation_unattended: float = 0.0  # all oxidation attended
    cooling: float = 3.0
    # Filtration & washing
    filtration_setup: float = 1.0
    filtration_unattended: float = 8.0
    wash_repulp: float = 1.0
    wash_filtration_unattended: float = 6.0
    # Charging/discharging
    charge_solids: float = 1.0
    fill_water: float = 0.5
    discharge_to_dryer: float = 1.0
    # Drying & wrap-up
    drying_setup: float = 1.0
    drying_unattended: float = 12.0
    packaging_qc: float = 1.0
    cleaning_changeover: float = 2.0


@dataclass
class RoleAllocs:
    supervisor_h_per_batch: float = 2.0
    qc_h_per_batch: float = 1.0
    maintenance_h_per_batch: float = 0.3
    eng_hse_h_per_batch: float = 0.5


@dataclass
class CampaignAllocs:
    supervisor_report_h: float = 24.0
    qc_report_h: float = 8.0
    eng_hse_report_h: float = 8.0


def batches_total(p: ProcessParams) -> int:
    return math.ceil(p.target_dry_kg / p.batch_dry_kg)


def slurry_mass_kg(p: ProcessParams) -> float:
    return p.batch_dry_kg / p.solids_mass_frac


def heatup_kwh(p: ProcessParams) -> float:
    m = slurry_mass_kg(p)
    energy_kj = m * p.cp_kj_per_kgk * (p.t_hot_c - p.t_cold_c)
    kwh = energy_kj / 3600
    return kwh / p.elec_efficiency * p.heat_overhead_factor


def cooling_kwh_th(p: ProcessParams) -> float:
    m = slurry_mass_kg(p)
    energy_kj = m * p.cp_kj_per_kgk * (p.t_hot_c - p.t_filter_c)
    return energy_kj / 3600


def oxidation_total_h(d: DurationsHours) -> float:
    return d.oxidation_attended + d.oxidation_unattended


def electricity_kwh(p: ProcessParams, d: DurationsHours, pow: EquipmentPowers, air: AirSystem) -> Dict[str, float]:
    oxid_total = oxidation_total_h(d)
    heatup_energy = heatup_kwh(p)
    heat_loss = pow.heat_loss_kw * oxid_total
    agitator = pow.agitator_kw * (d.heatup + oxid_total + d.cooling + d.charge_solids + d.fill_water + d.wash_repulp)
    circulation = pow.circulation_pump_kw * oxid_total
    air_comp = air.flow_m3ph * air.sec_kwh_per_m3 * oxid_total
    dryer = pow.dryer_kw * d.drying_unattended
    misc = pow.misc_kwht_per_batch
    total = heatup_energy + heat_loss + agitator + circulation + air_comp + dryer + misc
    return {
        "heatup": heatup_energy,
        "heat_loss": heat_loss,
        "agitator": agitator,
        "circulation_pump": circulation,
        "air_compressor": air_comp,
        "dryer": dryer,
        "misc": misc,
        "total": total,
    }


def manhours_operator(d: DurationsHours) -> float:
    return (
        2 * d.heatup +
        1 * d.oxidation_attended +
        1 * d.cooling +
        2 * d.filtration_setup +
        1 * d.drying_setup +
        1 * d.packaging_qc +
        2 * d.cleaning_changeover +
        1 * d.charge_solids +
        1 * d.fill_water +
        1 * d.discharge_to_dryer +
        1 * d.wash_repulp
    )


def costs_per_batch(p: ProcessParams, r: Rates, elec: Dict[str, float], mh_op: float, ra: RoleAllocs) -> Dict[str, float]:
    # Labor
    labor_operator = mh_op * r.operator
    labor_supervisor = ra.supervisor_h_per_batch * r.supervisor
    labor_qc = ra.qc_h_per_batch * r.qc
    labor_maintenance = ra.maintenance_h_per_batch * r.maintenance
    labor_eng_hse = ra.eng_hse_h_per_batch * r.eng_hse
    labor_total = sum([labor_operator, labor_supervisor, labor_qc, labor_maintenance, labor_eng_hse])

    # Utilities/media
    electricity_cost = elec["total"] * r.electricity_pln_per_kwh
    chilled_cost = cooling_kwh_th(p) * r.chilled_pln_per_kwhc if p.use_chiller else 0.0
    cooling_water_cost = p.cooling_water_m3_per_batch * r.cooling_water_pln_per_m3
    process_water_cost = p.process_water_m3_per_batch * r.process_water_pln_per_m3
    consumables_cost = p.consumables_pln_per_batch
    utilities_media_total = electricity_cost + chilled_cost + cooling_water_cost + process_water_cost + consumables_cost

    equipment_cost = 0.0

    base = labor_total + utilities_media_total + equipment_cost
    admin = p.admin_pct * (labor_total + equipment_cost)
    contingency_base = base if p.contingency_on_base_pre_admin else (base + admin)
    contingency = p.contingency_pct * contingency_base
    total = base + admin + contingency

    return {
        "labor_operator": labor_operator,
        "labor_supervisor": labor_supervisor,
        "labor_qc": labor_qc,
        "labor_maintenance": labor_maintenance,
        "labor_eng_hse": labor_eng_hse,
        "labor_total": labor_total,
        "electricity": electricity_cost,
        "chilled_energy": chilled_cost,
        "cooling_water": cooling_water_cost,
        "process_water": process_water_cost,
        "consumables": consumables_cost,
        "utilities_media_total": utilities_media_total,
        "equipment": equipment_cost,
        "admin": admin,
        "contingency": contingency,
        "total": total,
    }


def campaign_level_costs(r: Rates, p: ProcessParams, ca: CampaignAllocs) -> Dict[str, float]:
    oneoff_supervisor = ca.supervisor_report_h * r.supervisor
    oneoff_qc = ca.qc_report_h * r.qc
    oneoff_eng = ca.eng_hse_report_h * r.eng_hse
    oneoff_labor = oneoff_supervisor + oneoff_qc + oneoff_eng
    admin_on_oneoff = p.admin_pct * oneoff_labor
    contingency_base = oneoff_labor if p.contingency_on_base_pre_admin else (oneoff_labor + admin_on_oneoff)
    contingency_on_oneoff = p.contingency_pct * contingency_base
    total_oneoffs = oneoff_labor + admin_on_oneoff + contingency_on_oneoff
    return {
        "oneoff_supervisor": oneoff_supervisor,
        "oneoff_qc": oneoff_qc,
        "oneoff_eng_hse": oneoff_eng,
        "labor_oneoff_total": oneoff_labor,
        "admin_on_oneoff": admin_on_oneoff,
        "contingency_on_oneoff": contingency_on_oneoff,
        "total_oneoffs": total_oneoffs,
    }


def summary_dict(p: ProcessParams, d: DurationsHours, elec: Dict[str, float], costs: Dict[str, float], mh_op: float,
                 ra: RoleAllocs, batches: int, ca: CampaignAllocs, oneoffs: Dict[str, float]) -> Dict[str, Any]:
    total_mh = mh_op + ra.supervisor_h_per_batch + ra.qc_h_per_batch + ra.maintenance_h_per_batch + ra.eng_hse_h_per_batch
    per_batch_total = costs["total"]
    campaign_totals = {
        "labor": costs["labor_total"] * batches + oneoffs["labor_oneoff_total"],
        "utilities_media": costs["utilities_media_total"] * batches,
        "equipment": costs["equipment"] * batches,
        "admin": costs["admin"] * batches + oneoffs["admin_on_oneoff"],
        "contingency": costs["contingency"] * batches + oneoffs["contingency_on_oneoff"],
    }
    campaign_totals["total"] = sum(campaign_totals.values())

    return {
        "assumptions": {
            "batch_dry_kg": p.batch_dry_kg,
            "solids_mass_fraction": p.solids_mass_frac,
            "target_campaign_kg": p.target_dry_kg,
            "durations": {
                "charge_solids_h": d.charge_solids,
                "fill_water_h": d.fill_water,
                "heatup_h": d.heatup,
                "oxidation_attended_h": d.oxidation_attended,
                "oxidation_unattended_h": d.oxidation_unattended,
                "cooling_h": d.cooling,
                "filtration_setup_h": d.filtration_setup,
                "filtration_unattended_h": d.filtration_unattended,
                "wash_repulp_h": d.wash_repulp,
                "wash_filtration_unattended_h": d.wash_filtration_unattended,
                "discharge_to_dryer_h": d.discharge_to_dryer,
                "drying_setup_h": d.drying_setup,
                "drying_unattended_h": d.drying_unattended,
                "packaging_qc_h": d.packaging_qc,
                "cleaning_changeover_h": d.cleaning_changeover,
            },
            "use_chiller": p.use_chiller,
        },
        "derived": {
            "slurry_mass_kg": slurry_mass_kg(p),
            "batches_total": batches,
            "cooling_duty_kwh_th": cooling_kwh_th(p),
        },
        "manhours": {
            "operator": mh_op,
            "supervisor": ra.supervisor_h_per_batch,
            "qc": ra.qc_h_per_batch,
            "maintenance": ra.maintenance_h_per_batch,
            "eng_hse": ra.eng_hse_h_per_batch,
            "total": total_mh,
        },
        "utilities_per_batch": {
            "electricity_kwh": elec,
            "cooling_water_m3": p.cooling_water_m3_per_batch,
            "process_water_m3": p.process_water_m3_per_batch,
        },
        "costs_per_batch": costs,
        "campaign_oneoffs": oneoffs,
        "campaign_costs": campaign_totals,
    }


def generate_markdown_report(summary: Dict[str, Any], output_path: str) -> None:
    with open(output_path, 'w') as f:
        f.write("# ORLEN Lignin Oxidation Process Simulation Results\n\n")
        f.write("## Assumptions\n\n")
        f.write(f"- Batch size: **{summary['assumptions']['batch_dry_kg']:.1f} kg** (dry solids)\n")
        f.write(f"- Solids mass fraction: **{summary['assumptions']['solids_mass_fraction']:.2f}**\n")
        f.write(f"- Target campaign: **{summary['assumptions']['target_campaign_kg']:.1f} kg**\n")

        f.write("\n### Process Step Durations\n\n")
        f.write("| Step | Duration (h) | Crew |\n")
        f.write("|------|------------:|------:|\n")
        f.write(f"| Charge solids | {summary['assumptions']['durations']['charge_solids_h']:.1f} | 1 |\n")
        f.write(f"| Fill water | {summary['assumptions']['durations']['fill_water_h']:.1f} | 1 |\n")
        f.write(f"| Heat-up | {summary['assumptions']['durations']['heatup_h']:.1f} | 2 |\n")
        f.write(f"| Oxidation (attended) | {summary['assumptions']['durations']['oxidation_attended_h']:.1f} | 1 |\n")
        f.write(f"| Cooling | {summary['assumptions']['durations']['cooling_h']:.1f} | 1 |\n")
        f.write(f"| Filtration setup | {summary['assumptions']['durations']['filtration_setup_h']:.1f} | 2 |\n")
        f.write(f"| Filtration (unattended) | {summary['assumptions']['durations']['filtration_unattended_h']:.1f} | 0 |\n")
        f.write(f"| Wash repulpation | {summary['assumptions']['durations']['wash_repulp_h']:.1f} | 1 |\n")
        f.write(f"| Wash filtration (unattended) | {summary['assumptions']['durations']['wash_filtration_unattended_h']:.1f} | 0 |\n")
        f.write(f"| Discharge to dryer | {summary['assumptions']['durations']['discharge_to_dryer_h']:.1f} | 1 |\n")
        f.write(f"| Drying setup | {summary['assumptions']['durations']['drying_setup_h']:.1f} | 1 |\n")
        f.write(f"| Drying (unattended) | {summary['assumptions']['durations']['drying_unattended_h']:.1f} | 0 |\n")
        f.write(f"| Packaging & QC | {summary['assumptions']['durations']['packaging_qc_h']:.1f} | 1 |\n")
        f.write(f"| Cleaning/changeover | {summary['assumptions']['durations']['cleaning_changeover_h']:.1f} | 2 |\n")

        f.write("\n## Derived Values\n\n")
        f.write(f"- Slurry mass: **{summary['derived']['slurry_mass_kg']:.1f} kg** per batch\n")
        f.write(f"- Total batches: **{summary['derived']['batches_total']}**\n")
        f.write(f"- Cooling duty: **{summary['derived']['cooling_duty_kwh_th']:.1f} kWh_th** per batch\n")
        if not summary['assumptions']['use_chiller']:
            f.write("\n> **Note:** Using cooling water instead of chiller. Avoid double-counting cooling costs.\n")

        f.write("\n## Man-Hours\n\n")
        f.write("| Role | Hours/Batch |\n")
        f.write("|------|------------:|\n")
        f.write(f"| Operator | {summary['manhours']['operator']:.1f} |\n")
        f.write(f"| Supervisor | {summary['manhours']['supervisor']:.1f} |\n")
        f.write(f"| QC/Lab | {summary['manhours']['qc']:.1f} |\n")
        f.write(f"| Maintenance | {summary['manhours']['maintenance']:.1f} |\n")
        f.write(f"| Engineering/HSE | {summary['manhours']['eng_hse']:.1f} |\n")
        f.write(f"| **Total** | **{summary['manhours']['total']:.1f}** |\n")

        f.write("\n## Utilities Per Batch\n\n")
        f.write("### Electricity Breakdown\n\n")
        f.write("| Consumer | kWh | PLN |\n")
        f.write("|----------|----:|----:|\n")
        elec = summary['utilities_per_batch']['electricity_kwh']
        rate = 1.50
        f.write(f"| Heat-up | {elec['heatup']:.2f} | {elec['heatup'] * rate:.2f} |\n")
        f.write(f"| Heat loss | {elec['heat_loss']:.2f} | {elec['heat_loss'] * rate:.2f} |\n")
        f.write(f"| Agitator | {elec['agitator']:.2f} | {elec['agitator'] * rate:.2f} |\n")
        f.write(f"| Circulation pump | {elec['circulation_pump']:.2f} | {elec['circulation_pump'] * rate:.2f} |\n")
        f.write(f"| Air compressor | {elec['air_compressor']:.2f} | {elec['air_compressor'] * rate:.2f} |\n")
        f.write(f"| Dryer | {elec['dryer']:.2f} | {elec['dryer'] * rate:.2f} |\n")
        f.write(f"| Misc | {elec['misc']:.2f} | {elec['misc'] * rate:.2f} |\n")
        f.write(f"| **Total** | **{elec['total']:.2f}** | **{elec['total'] * rate:.2f}** |\n")

        f.write("\n### Other Utilities\n\n")
        f.write(f"- Cooling water: **{summary['utilities_per_batch']['cooling_water_m3']:.1f} m³** per batch\n")
        f.write(f"- Process water: **{summary['utilities_per_batch']['process_water_m3']:.1f} m³** per batch\n")

        f.write("\n## Costs\n\n")
        f.write("### Per Batch (PLN)\n\n")
        f.write("| Category | Cost (PLN) |\n")
        f.write("|----------|------------:|\n")
        f.write(f"| Labor - Operator | {summary['costs_per_batch']['labor_operator']:.2f} |\n")
        f.write(f"| Labor - Supervisor | {summary['costs_per_batch']['labor_supervisor']:.2f} |\n")
        f.write(f"| Labor - QC/Lab | {summary['costs_per_batch']['labor_qc']:.2f} |\n")
        f.write(f"| Labor - Maintenance | {summary['costs_per_batch']['labor_maintenance']:.2f} |\n")
        f.write(f"| Labor - Eng/HSE | {summary['costs_per_batch']['labor_eng_hse']:.2f} |\n")
        f.write(f"| **Labor Subtotal** | **{summary['costs_per_batch']['labor_total']:.2f}** |\n")
        f.write(f"| Electricity | {summary['costs_per_batch']['electricity']:.2f} |\n")
        if summary['assumptions']['use_chiller']:
            f.write(f"| Chilled energy | {summary['costs_per_batch']['chilled_energy']:.2f} |\n")
        f.write(f"| Cooling water | {summary['costs_per_batch']['cooling_water']:.2f} |\n")
        f.write(f"| Process water | {summary['costs_per_batch']['process_water']:.2f} |\n")
        f.write(f"| Consumables | {summary['costs_per_batch']['consumables']:.2f} |\n")
        f.write(f"| **Utilities/Media Subtotal** | **{summary['costs_per_batch']['utilities_media_total']:.2f}** |\n")
        f.write(f"| Equipment | {summary['costs_per_batch']['equipment']:.2f} |\n")
        f.write(f"| Admin (10%) | {summary['costs_per_batch']['admin']:.2f} |\n")
        f.write(f"| Contingency (15%) | {summary['costs_per_batch']['contingency']:.2f} |\n")
        f.write(f"| **Total Per Batch** | **{summary['costs_per_batch']['total']:.2f}** |\n")

        f.write("\n### Campaign One-off Labor (PLN)\n\n")
        f.write("| Role | Hours | Cost |\n")
        f.write("|------|------:|-----:|\n")
        f.write(f"| Supervisor (report/analysis) | 24.0 | {summary['campaign_oneoffs']['oneoff_supervisor']:.2f} |\n")
        f.write(f"| QC/Lab (report/validation) | 8.0 | {summary['campaign_oneoffs']['oneoff_qc']:.2f} |\n")
        f.write(f"| Eng/HSE (review) | 8.0 | {summary['campaign_oneoffs']['oneoff_eng_hse']:.2f} |\n")
        f.write(f"| **Labor one-offs subtotal** |  | **{summary['campaign_oneoffs']['labor_oneoff_total']:.2f}** |\n")
        f.write(f"| Admin on one-offs (10%) |  | {summary['campaign_oneoffs']['admin_on_oneoff']:.2f} |\n")
        f.write(f"| Contingency on one-offs (15%) |  | {summary['campaign_oneoffs']['contingency_on_oneoff']:.2f} |\n")
        f.write(f"| **One-offs total** |  | **{summary['campaign_oneoffs']['total_oneoffs']:.2f}** |\n")

        f.write("\n### Campaign Total (PLN)\n\n")
        f.write("| Category | Cost (PLN) |\n")
        f.write("|----------|------------:|\n")
        f.write(f"| Labor | {summary['campaign_costs']['labor']:.2f} |\n")
        f.write(f"| Utilities/Media | {summary['campaign_costs']['utilities_media']:.2f} |\n")
        f.write(f"| Equipment | {summary['campaign_costs']['equipment']:.2f} |\n")
        f.write(f"| Admin | {summary['campaign_costs']['admin']:.2f} |\n")
        f.write(f"| Contingency | {summary['campaign_costs']['contingency']:.2f} |\n")
        f.write(f"| **Campaign Total** | **{summary['campaign_costs']['total']:.2f}** |\n")

        f.write("\n---\n\n")
        f.write("Generated by process_simulation.py on ORLEN_2025-08_RFI-oxidation project\n")


if __name__ == "__main__":
    rates = Rates()
    powers = EquipmentPowers()
    air = AirSystem()
    params = ProcessParams()
    durs = DurationsHours()
    role_allocs = RoleAllocs()
    camp_allocs = CampaignAllocs()

    batches = batches_total(params)
    slurry_kg = slurry_mass_kg(params)

    elec = electricity_kwh(params, durs, powers, air)
    mh_op = manhours_operator(durs)
    costs = costs_per_batch(params, rates, elec, mh_op, role_allocs)
    oneoffs = campaign_level_costs(rates, params, camp_allocs)
    summary = summary_dict(params, durs, elec, costs, mh_op, role_allocs, batches, camp_allocs, oneoffs)

    print("\n=== ORLEN Lignin Oxidation Process Simulation ===\n")
    print(f"Batch size: {params.batch_dry_kg:.1f} kg dry solids ({slurry_kg:.1f} kg slurry)")
    print(f"Campaign target: {params.target_dry_kg:.1f} kg → {batches} batches\n")

    print("Man-hours per batch:")
    print(f"  Operator: {mh_op:.1f} h")
    total_mh = mh_op + role_allocs.supervisor_h_per_batch + role_allocs.qc_h_per_batch + role_allocs.maintenance_h_per_batch + role_allocs.eng_hse_h_per_batch
    print(f"  Other roles: {total_mh - mh_op:.1f} h")
    print(f"  Total: {total_mh:.1f} h\n")

    print("Electricity per batch:")
    print(f"  Heat-up: {elec['heatup']:.2f} kWh")
    print(f"  Heat loss (hold): {elec['heat_loss']:.2f} kWh")
    print(f"  Agitator: {elec['agitator']:.2f} kWh")
    print(f"  Circulation pump: {elec['circulation_pump']:.2f} kWh")
    print(f"  Air compressor: {elec['air_compressor']:.2f} kWh")
    print(f"  Dryer: {elec['dryer']:.2f} kWh")
    print(f"  Misc: {elec['misc']:.2f} kWh")
    print(f"  Total: {elec['total']:.2f} kWh\n")

    # Per-batch costs
    print("Costs per batch (PLN):")
    print(f"  Labor: {costs['labor_total']:.2f}")
    print(f"  Utilities/Media: {costs['utilities_media_total']:.2f}")
    print(f"  Equipment: {costs['equipment']:.2f}")
    print(f"  Admin: {costs['admin']:.2f}")
    print(f"  Contingency: {costs['contingency']:.2f}")
    print(f"  Total per batch: {costs['total']:.2f}\n")

    # Campaign totals including one-offs
    print("Campaign one-offs (PLN):")
    print(f"  Supervisor: {oneoffs['oneoff_supervisor']:.2f}")
    print(f"  QC/Lab: {oneoffs['oneoff_qc']:.2f}")
    print(f"  Eng/HSE: {oneoffs['oneoff_eng_hse']:.2f}")
    print(f"  Admin on one-offs: {oneoffs['admin_on_oneoff']:.2f}")
    print(f"  Contingency on one-offs: {oneoffs['contingency_on_oneoff']:.2f}")
    print(f"  One-offs total: {oneoffs['total_oneoffs']:.2f}\n")

    campaign_total = summary['campaign_costs']['total']
    print(f"Campaign total (incl. one-offs): {campaign_total:.2f} PLN\n")

    # Ensure output directory exists
    output_dir = Path("projects/ORLEN_2025-08_RFI-oxidation/outputs")
    os.makedirs(output_dir, exist_ok=True)

    # Write JSON and MD
    json_path = output_dir / "simulation_results.json"
    with open(json_path, 'w') as f:
        json.dump(summary, f, indent=2)
    md_path = output_dir / "simulation_results.md"
    generate_markdown_report(summary, str(md_path))
    print(f"JSON → {json_path}\nMD   → {md_path}")