"""
hydraulic_core.py

Shared deterministic hydraulic teaching functions used by both:
- server.py    -> MCP surface
- rest_api.py  -> REST / OpenAPI surface

Units:
- MKS display: L/s, m, mm
- FPS display: gpm, ft, in
Internal calculations:
- flow: m3/s
- length/head: m
- diameter: m
"""

from __future__ import annotations

import math
import time
import uuid
from dataclasses import dataclass, asdict
from typing import Any, Dict, List, Optional

HW_EXP = 1.852

M3S_TO_LPS = 1000.0
LPS_TO_M3S = 0.001
M3S_TO_GPM = 15850.323141489
GPM_TO_M3S = 1.0 / M3S_TO_GPM
M_TO_FT = 3.280839895
FT_TO_M = 1.0 / M_TO_FT
M_TO_IN = 39.37007874
IN_TO_M = 1.0 / M_TO_IN
MM_TO_M = 0.001
M_HEAD_TO_PSI = 1.421970206
PSI_TO_M_HEAD = 1.0 / M_HEAD_TO_PSI


def unit_mode(unit_system: str = "MKS / L/s") -> str:
    return "FPS" if isinstance(unit_system, str) and "FPS" in unit_system else "MKS"


def labels(unit_system: str = "MKS / L/s") -> Dict[str, str]:
    if unit_mode(unit_system) == "FPS":
        return {"flow": "gpm", "length": "ft", "diameter": "in", "head": "ft", "pressure": "psi"}
    return {"flow": "L/s", "length": "m", "diameter": "mm", "head": "m", "pressure": "m"}


def flow_to_m3s(x: float, unit_system: str = "MKS / L/s") -> float:
    return float(x) * (GPM_TO_M3S if unit_mode(unit_system) == "FPS" else LPS_TO_M3S)


def m3s_to_flow(x: float, unit_system: str = "MKS / L/s") -> float:
    return float(x) * (M3S_TO_GPM if unit_mode(unit_system) == "FPS" else M3S_TO_LPS)


def length_to_m(x: float, unit_system: str = "MKS / L/s") -> float:
    return float(x) * (FT_TO_M if unit_mode(unit_system) == "FPS" else 1.0)


def m_to_length(x: float, unit_system: str = "MKS / L/s") -> float:
    return float(x) * (M_TO_FT if unit_mode(unit_system) == "FPS" else 1.0)


def head_to_m(x: float, unit_system: str = "MKS / L/s") -> float:
    return length_to_m(x, unit_system)


def m_to_head(x: float, unit_system: str = "MKS / L/s") -> float:
    return m_to_length(x, unit_system)


def diameter_to_m(x: float, unit_system: str = "MKS / L/s") -> float:
    return float(x) * (IN_TO_M if unit_mode(unit_system) == "FPS" else MM_TO_M)


def pressure_head_to_display(x_m: float, unit_system: str = "MKS / L/s") -> float:
    return float(x_m) * (M_HEAD_TO_PSI if unit_mode(unit_system) == "FPS" else 1.0)


def hw_resistance(length_m: float, diameter_m: float, c_hw: float) -> float:
    if length_m <= 0 or diameter_m <= 0 or c_hw <= 0:
        raise ValueError("Pipe length, diameter, and Hazen-Williams C must be positive.")
    return 10.67 * length_m / ((float(c_hw) ** HW_EXP) * (diameter_m ** 4.871))


def signed_headloss(q_m3s: float, r: float) -> float:
    if abs(q_m3s) < 1e-15:
        return 0.0
    return math.copysign(r * abs(q_m3s) ** HW_EXP, q_m3s)


def flow_from_headloss(delta_h_m: float, r: float) -> float:
    if abs(delta_h_m) < 1e-15:
        return 0.0
    return math.copysign((abs(delta_h_m) / r) ** (1.0 / HW_EXP), delta_h_m)


@dataclass
class NetworkSession:
    session_id: str
    created_at: float
    last_used: float
    inp_text: str
    title: str
    units: str
    counts: Dict[str, int]


SESSIONS: Dict[str, NetworkSession] = {}
SESSION_TTL_SECONDS = 3600


def _expire_sessions() -> None:
    now = time.time()
    expired = [sid for sid, s in SESSIONS.items() if now - s.last_used > SESSION_TTL_SECONDS]
    for sid in expired:
        SESSIONS.pop(sid, None)


def _section_counts(inp_text: str) -> Dict[str, int]:
    counts: Dict[str, int] = {}
    current = None
    for raw in inp_text.splitlines():
        line = raw.strip()
        if not line or line.startswith(";"):
            continue
        if line.startswith("[") and line.endswith("]"):
            current = line.upper()
            counts[current] = 0
        elif current:
            counts[current] += 1
    return counts


def _detect_units(inp_text: str) -> str:
    for raw in inp_text.splitlines():
        line = raw.strip().upper()
        if line.startswith("UNITS"):
            if "GPM" in line:
                return "GPM"
            if "LPS" in line:
                return "LPS"
            return line.split()[-1]
    return "UNKNOWN"


def load_network(inp_text: str, title: str = "uploaded_network") -> Dict[str, Any]:
    """Create a session from EPANET INP text."""
    _expire_sessions()
    sid = str(uuid.uuid4())
    session = NetworkSession(
        session_id=sid,
        created_at=time.time(),
        last_used=time.time(),
        inp_text=str(inp_text),
        title=title,
        units=_detect_units(inp_text),
        counts=_section_counts(inp_text),
    )
    SESSIONS[sid] = session
    return {
        "session_id": sid,
        "title": title,
        "units": session.units,
        "counts": session.counts,
        "message": "Network session loaded. Use this session_id with analysis tools.",
    }


def get_session(session_id: str) -> NetworkSession:
    _expire_sessions()
    if session_id not in SESSIONS:
        raise KeyError(f"Unknown or expired session_id: {session_id}")
    SESSIONS[session_id].last_used = time.time()
    return SESSIONS[session_id]


def close_session(session_id: str) -> Dict[str, Any]:
    existed = session_id in SESSIONS
    SESSIONS.pop(session_id, None)
    return {"session_id": session_id, "closed": existed}


def network_summary(session_id: str) -> Dict[str, Any]:
    s = get_session(session_id)
    return {
        "session_id": s.session_id,
        "title": s.title,
        "units": s.units,
        "counts": s.counts,
        "age_seconds": round(time.time() - s.created_at, 3),
        "note": "This summary is parsed from INP text. Full EPANET/WNTR simulation can be wired in a later engine layer.",
    }


def solve_single_pipe(unit_system: str, length: float, diameter: float, c_hw: float, target_headloss: float,
                      initial_flow: float, max_iter: int = 25, tolerance: float = 1e-4) -> Dict[str, Any]:
    L = length_to_m(length, unit_system)
    D = diameter_to_m(diameter, unit_system)
    target = head_to_m(target_headloss, unit_system)
    q = abs(flow_to_m3s(initial_flow, unit_system))
    tol = head_to_m(tolerance, unit_system)
    r = hw_resistance(L, D, c_hw)
    rows = []
    for i in range(1, int(max_iter) + 1):
        hf = r * abs(q) ** HW_EXP
        residual = hf - target
        derivative = HW_EXP * r * max(abs(q), 1e-12) ** (HW_EXP - 1.0)
        correction = residual / derivative if derivative else 0.0
        q_new = max(q - correction, 1e-12)
        rows.append({
            "iteration": i,
            "flow": m3s_to_flow(q, unit_system),
            "headloss": m_to_head(hf, unit_system),
            "residual": m_to_head(residual, unit_system),
            "correction_flow": m3s_to_flow(correction, unit_system),
            "updated_flow": m3s_to_flow(q_new, unit_system),
        })
        q = q_new
        if abs(residual) <= tol:
            break
    lab = labels(unit_system)
    return {
        "unit_system": unit_system,
        "units": lab,
        "final_flow": m3s_to_flow(q, unit_system),
        "final_headloss": m_to_head(r * abs(q) ** HW_EXP, unit_system),
        "iterations": rows,
    }


def solve_hardy_cross_loop(unit_system: str, flows: List[float], lengths: List[float], diameters: List[float],
                           c_values: List[float], max_iter: int = 25, tolerance: float = 1e-5) -> Dict[str, Any]:
    n = len(flows)
    if not (len(lengths) == len(diameters) == len(c_values) == n):
        raise ValueError("flows, lengths, diameters, and c_values must have equal length.")
    q = [flow_to_m3s(v, unit_system) for v in flows]
    r = [hw_resistance(length_to_m(lengths[i], unit_system), diameter_to_m(diameters[i], unit_system), c_values[i]) for i in range(n)]
    tol = abs(flow_to_m3s(tolerance, unit_system))
    summary = []
    details = []
    for it in range(1, int(max_iter) + 1):
        hs = [signed_headloss(q[i], r[i]) for i in range(n)]
        denom_terms = [HW_EXP * r[i] * max(abs(q[i]), 1e-12) ** (HW_EXP - 1.0) for i in range(n)]
        sum_h = sum(hs)
        denom = sum(denom_terms)
        dq = -sum_h / denom if denom else 0.0
        q_after = [v + dq for v in q]
        summary.append({
            "iteration": it,
            "sum_headloss": m_to_head(sum_h, unit_system),
            "correction": m3s_to_flow(dq, unit_system),
            "abs_correction": abs(m3s_to_flow(dq, unit_system)),
        })
        for i in range(n):
            details.append({
                "iteration": it,
                "pipe": f"P{i+1}",
                "flow_before": m3s_to_flow(q[i], unit_system),
                "signed_headloss": m_to_head(hs[i], unit_system),
                "flow_after": m3s_to_flow(q_after[i], unit_system),
                "direction": "same as assumed" if q_after[i] >= 0 else "opposite to assumed",
            })
        q = q_after
        if abs(dq) <= tol:
            break
    return {
        "unit_system": unit_system,
        "units": labels(unit_system),
        "final_flows": [{"pipe": f"P{i+1}", "flow": m3s_to_flow(q[i], unit_system), "direction": "same as assumed" if q[i] >= 0 else "opposite to assumed"} for i in range(n)],
        "iterations": summary,
        "details": details,
    }


def solve_two_loop_hardy_cross(unit_system: str, flows: List[float], common_length: float, common_diameter: float,
                               c_hw: float = 120.0, max_iter: int = 25, tolerance: float = 1e-5) -> Dict[str, Any]:
    if len(flows) != 7:
        raise ValueError("Two-loop teaching network requires 7 initial pipe flows.")
    q = [flow_to_m3s(v, unit_system) for v in flows]
    r = [hw_resistance(length_to_m(common_length, unit_system), diameter_to_m(common_diameter, unit_system), c_hw) for _ in range(7)]
    loop1, s1 = [0, 1, 2, 3], [1, 1, 1, 1]
    loop2, s2 = [4, 5, 6, 1], [1, 1, 1, -1]
    tol = abs(flow_to_m3s(tolerance, unit_system))
    iterations = []
    for it in range(1, int(max_iter) + 1):
        def corr(indices, signs):
            ql = [signs[j] * q[idx] for j, idx in enumerate(indices)]
            hl = [signed_headloss(ql[j], r[indices[j]]) for j in range(len(indices))]
            den = sum(HW_EXP * r[indices[j]] * max(abs(ql[j]), 1e-12) ** (HW_EXP - 1.0) for j in range(len(indices)))
            return sum(hl), (-sum(hl) / den if den else 0.0)
        sh1, dq1 = corr(loop1, s1)
        sh2, dq2 = corr(loop2, s2)
        for j, idx in enumerate(loop1):
            q[idx] += s1[j] * dq1
        for j, idx in enumerate(loop2):
            q[idx] += s2[j] * dq2
        iterations.append({
            "iteration": it,
            "loop1_sum_headloss": m_to_head(sh1, unit_system),
            "loop1_correction": m3s_to_flow(dq1, unit_system),
            "loop2_sum_headloss": m_to_head(sh2, unit_system),
            "loop2_correction": m3s_to_flow(dq2, unit_system),
            "max_abs_correction": max(abs(m3s_to_flow(dq1, unit_system)), abs(m3s_to_flow(dq2, unit_system))),
        })
        if max(abs(dq1), abs(dq2)) <= tol:
            break
    return {
        "unit_system": unit_system,
        "units": labels(unit_system),
        "final_flows": [{"pipe": f"P{i+1}", "flow": m3s_to_flow(q[i], unit_system), "shared_pipe": i == 1, "direction": "same as assumed" if q[i] >= 0 else "opposite to assumed"} for i in range(7)],
        "iterations": iterations,
    }


def solve_three_reservoir(unit_system: str, reservoir_heads: List[float], demand: float, initial_head: float,
                          lengths: List[float], diameters: List[float], c_values: List[float],
                          max_iter: int = 25, tolerance: float = 1e-4) -> Dict[str, Any]:
    H = [head_to_m(v, unit_system) for v in reservoir_heads]
    HJ = head_to_m(initial_head, unit_system)
    demand_m3s = flow_to_m3s(demand, unit_system)
    tol = abs(flow_to_m3s(tolerance, unit_system))
    r = [hw_resistance(length_to_m(lengths[i], unit_system), diameter_to_m(diameters[i], unit_system), c_values[i]) for i in range(3)]
    rows = []
    def residual_at(hj):
        qs = [flow_from_headloss(H[i] - hj, r[i]) for i in range(3)]
        return sum(qs) - demand_m3s, qs
    for it in range(1, int(max_iter) + 1):
        res, qs = residual_at(HJ)
        dH = 1e-4
        res_plus, _ = residual_at(HJ + dH)
        deriv = (res_plus - res) / dH
        corr = -res / deriv if abs(deriv) > 1e-12 else 0.0
        corr = max(min(corr, 20.0), -20.0)
        rows.append({
            "iteration": it,
            "junction_head": m_to_head(HJ, unit_system),
            "q1": m3s_to_flow(qs[0], unit_system),
            "q2": m3s_to_flow(qs[1], unit_system),
            "q3": m3s_to_flow(qs[2], unit_system),
            "continuity_residual": m3s_to_flow(res, unit_system),
            "head_correction": m_to_head(corr, unit_system),
            "updated_junction_head": m_to_head(HJ + corr, unit_system),
        })
        HJ += corr
        if abs(res) <= tol:
            break
    res, qs = residual_at(HJ)
    return {
        "unit_system": unit_system,
        "units": labels(unit_system),
        "junction_head": m_to_head(HJ, unit_system),
        "flows": [{"pipe": f"P{i+1}", "flow": m3s_to_flow(qs[i], unit_system), "meaning": "reservoir supplies junction" if qs[i] >= 0 else "junction supplies reservoir"} for i in range(3)],
        "continuity_residual": m3s_to_flow(res, unit_system),
        "iterations": rows,
    }


def solve_pdd_demand(unit_system: str, required_demand: float, available_pressure: float,
                     minimum_pressure: float, required_pressure: float, exponent: float = 0.5) -> Dict[str, Any]:
    req = flow_to_m3s(required_demand, unit_system)
    p = head_to_m(available_pressure, unit_system)
    pmin = head_to_m(minimum_pressure, unit_system)
    preq = head_to_m(required_pressure, unit_system)
    if p <= pmin:
        ratio = 0.0
    elif p >= preq:
        ratio = 1.0
    else:
        ratio = ((p - pmin) / max(preq - pmin, 1e-12)) ** float(exponent)
    delivered = req * ratio
    return {
        "unit_system": unit_system,
        "units": labels(unit_system),
        "satisfaction_ratio": ratio,
        "required_demand": required_demand,
        "delivered_demand": m3s_to_flow(delivered, unit_system),
        "unserved_demand": m3s_to_flow(req - delivered, unit_system),
    }


def simulate_tank_eps(unit_system: str, diameter: float, initial_level: float, min_level: float, max_level: float,
                      timestep_hr: float, inflows: List[float], outflows: List[float]) -> Dict[str, Any]:
    D = length_to_m(diameter, unit_system)
    level = length_to_m(initial_level, unit_system)
    min_l = length_to_m(min_level, unit_system)
    max_l = length_to_m(max_level, unit_system)
    area = math.pi * D**2 / 4
    n = max(len(inflows), len(outflows))
    rows = []
    for i in range(n):
        qin = flow_to_m3s(inflows[i % len(inflows)], unit_system)
        qout = flow_to_m3s(outflows[i % len(outflows)], unit_system)
        start = level
        end_unclipped = start + (qin - qout) * float(timestep_hr) * 3600.0 / area
        end = min(max(end_unclipped, min_l), max_l)
        rows.append({
            "step": i + 1,
            "start_level": m_to_head(start, unit_system),
            "end_level": m_to_head(end, unit_system),
            "inflow": m3s_to_flow(qin, unit_system),
            "outflow": m3s_to_flow(qout, unit_system),
            "net_flow": m3s_to_flow(qin - qout, unit_system),
            "clipped": abs(end - end_unclipped) > 1e-9,
        })
        level = end
    return {"unit_system": unit_system, "units": labels(unit_system), "final_level": m_to_head(level, unit_system), "time_series": rows}


def evaluate_valve_behavior(unit_system: str, valve_type: str, upstream_head: float, setting: float,
                            flow: float, diameter: float, minor_loss_k: float = 0.0) -> Dict[str, Any]:
    H_up = head_to_m(upstream_head, unit_system)
    H_set = head_to_m(setting, unit_system)
    q = flow_to_m3s(flow, unit_system)
    D = diameter_to_m(diameter, unit_system)
    area = math.pi * D**2 / 4
    v = q / area if area else 0.0
    h_minor = float(minor_loss_k) * v**2 / (2 * 9.80665)
    vt = valve_type.upper()
    if vt == "PRV":
        H_down = min(H_up - h_minor, H_set)
        status = "active/throttling" if H_up > H_set else "open"
    elif vt == "PSV":
        H_down = H_up - h_minor
        status = "active/sustaining upstream" if H_up <= H_set else "open"
    elif vt == "FCV":
        H_down = H_up - h_minor
        status = "active/flow controlled if hydraulically feasible"
    elif vt == "TCV":
        H_down = H_up - h_minor
        status = "active/minor-loss control"
    else:
        status = "open" if H_up > H_set else "closed/reverse prevented"
        if "closed" in status:
            q = 0.0
            H_down = H_set
        else:
            H_down = H_up - h_minor
    return {
        "unit_system": unit_system,
        "units": labels(unit_system),
        "valve_type": valve_type,
        "status": status,
        "upstream_head": upstream_head,
        "downstream_head_estimate": m_to_head(H_down, unit_system),
        "headloss": m_to_head(H_up - H_down, unit_system),
        "flow": m3s_to_flow(q, unit_system),
        "teaching_note": "This isolates valve logic. Full network behavior should be validated in EPANET/WNTR.",
    }


def solve_pump_operating_point(unit_system: str, shutoff_head: float, design_flow: float, static_head: float,
                               pump_curve_k: float, system_curve_k: float) -> Dict[str, Any]:
    qmax = max(float(design_flow) * 1.8, 1e-9)
    q_values = [qmax * i / 200 for i in range(201)]
    pump_h = [float(shutoff_head) - float(pump_curve_k) * q**2 for q in q_values]
    sys_h = [float(static_head) + float(system_curve_k) * q**2 for q in q_values]
    idx = min(range(len(q_values)), key=lambda i: abs(pump_h[i] - sys_h[i]))
    return {
        "unit_system": unit_system,
        "units": labels(unit_system),
        "operating_flow": q_values[idx],
        "operating_head": pump_h[idx],
        "curve_table": [{"flow": q_values[i], "pump_head": pump_h[i], "system_head": sys_h[i]} for i in range(0, len(q_values), 10)],
    }


def pressure_zone_analysis(unit_system: str, source_head: float, prv_setting: float, node_elevations: List[float],
                           demand_multiplier: float = 1.0, min_pressure: float = 14.0, max_pressure: float = 56.0) -> Dict[str, Any]:
    H_source = head_to_m(source_head, unit_system)
    H_zone = min(H_source, head_to_m(prv_setting, unit_system))
    min_p_m = head_to_m(min_pressure, unit_system)
    max_p_m = head_to_m(max_pressure, unit_system)
    rows = []
    for i, elev in enumerate(node_elevations, start=1):
        e_m = length_to_m(elev, unit_system)
        p_head = H_zone - e_m
        if p_head < min_p_m:
            cat = "low pressure"
        elif p_head > max_p_m:
            cat = "high pressure / leakage risk"
        else:
            cat = "acceptable"
        rows.append({"node": f"J{i}", "elevation": elev, "pressure": pressure_head_to_display(p_head, unit_system), "category": cat})
    return {
        "unit_system": unit_system,
        "units": labels(unit_system),
        "zone_hgl": m_to_head(H_zone, unit_system),
        "prv_status": "active" if H_source > H_zone else "open/source limited",
        "nodes": rows,
        "teaching_note": "Pressure zone behavior depends on source head, control setting, elevation, and demands.",
    }


def leakage_nrw_analysis(unit_system: str, average_pressure: float, authorized_demand: float,
                         leakage_coefficient: float, pressure_exponent: float = 1.0) -> Dict[str, Any]:
    p = max(float(average_pressure), 0.0)
    leakage = float(leakage_coefficient) * (p ** float(pressure_exponent))
    total = float(authorized_demand) + leakage
    return {
        "unit_system": unit_system,
        "units": labels(unit_system),
        "authorized_demand": authorized_demand,
        "leakage_flow": leakage,
        "zone_inflow": total,
        "nrw_percent": 100.0 * leakage / total if total > 0 else 0.0,
        "teaching_note": "EPANET emitter nodes can be used to validate pressure-dependent leakage behavior.",
    }


def water_age_analysis(unit_system: str, pipe_volume: float, tank_volume: float, demand: float,
                       dead_end_factor: float = 2.0) -> Dict[str, Any]:
    # display volumes are assumed m3 in MKS and gallons in FPS
    if unit_mode(unit_system) == "FPS":
        volume_m3 = (float(pipe_volume) + float(tank_volume)) / 264.1720524
    else:
        volume_m3 = float(pipe_volume) + float(tank_volume)
    q = max(flow_to_m3s(demand, unit_system), 1e-12)
    base_age_hr = volume_m3 / q / 3600.0
    return {
        "unit_system": unit_system,
        "units": labels(unit_system),
        "average_age_hours": base_age_hr,
        "dead_end_age_hours": base_age_hr * float(dead_end_factor),
        "teaching_note": "Low flow, large storage, and dead ends increase water age. Validate with EPANET AGE analysis.",
    }


def chlorine_decay_analysis(initial_chlorine_mg_l: float, bulk_decay_per_day: float, travel_time_hours: float,
                            wall_decay_factor: float = 0.0) -> Dict[str, Any]:
    k_hr = (float(bulk_decay_per_day) + float(wall_decay_factor)) / 24.0
    c = float(initial_chlorine_mg_l) * math.exp(-k_hr * float(travel_time_hours))
    return {
        "initial_chlorine_mg_l": initial_chlorine_mg_l,
        "travel_time_hours": travel_time_hours,
        "final_chlorine_mg_l": c,
        "decay_fraction": 1.0 - c / max(float(initial_chlorine_mg_l), 1e-12),
        "teaching_note": "This first-order teaching model can be compared with EPANET chemical analysis.",
    }


def generate_epanet_validation_inp(case: str = "three_reservoir_mks") -> Dict[str, Any]:
    case = str(case).lower()
    if "fps" in case:
        text = """[TITLE]
Three-Reservoir Validation Model - FPS / GPM
[OPTIONS]
UNITS GPM
HEADLOSS H-W
[JUNCTIONS]
J 328.084 317.006
[RESERVOIRS]
R1 459.318
R2 393.701
R3 344.488
[PIPES]
P1 R1 J 1640.42 11.811 120 0 Open
P2 R2 J 2296.59 9.843 120 0 Open
P3 R3 J 1968.50 9.843 120 0 Open
[END]
"""
    else:
        text = """[TITLE]
Three-Reservoir Validation Model - MKS / LPS
[OPTIONS]
UNITS LPS
HEADLOSS H-W
[JUNCTIONS]
J 100 20
[RESERVOIRS]
R1 140
R2 120
R3 105
[PIPES]
P1 R1 J 500 300 120 0 Open
P2 R2 J 700 250 120 0 Open
P3 R3 J 600 250 120 0 Open
[END]
"""
    return {"case": case, "filename": f"{case}.inp", "inp_text": text}