from __future__ import annotations

from dataclasses import dataclass
from typing import List, Optional, Tuple
import numpy as np
from src.common.config import Config


@dataclass
class EncodedSchedule:
    """正演代理模型使用的固定长度流量制度表示。"""
    x_sched: np.ndarray
    x_sec: np.ndarray


def canonicalize_schedule(cfg: Config, timeQ: List[float], q: List[float]) -> Tuple[np.ndarray, np.ndarray, bool]:
    """清洗原始流量制度，同时保留求解器看到的主要流量变化形态。"""
    max_points = int(cfg.get("schedule", "max_points", default=512))
    dt = np.asarray(list(map(float, timeQ)), dtype=np.float64).reshape(-1)[:max_points]
    qq = np.asarray(list(map(float, q)), dtype=np.float64).reshape(-1)[:max_points]

    dt = np.maximum(dt, 1e-12)
    qq = np.maximum(qq, 0.0)

    cf = cfg.raw["schedule"].get("canonicalize_for_model", {}) or {}
    q_thr = float(cf.get("q_thr", 1e-6))
    merge_same_q = bool(cf.get("merge_same_q", True))
    merge_rel_tol = float(cf.get("merge_rel_tol", 1e-4))
    remove_shutin = bool(cf.get("remove_shutin", False))

    has_shutin = bool(np.any(qq <= q_thr))

    if merge_same_q and dt.size >= 2:
        # 合并相邻且流量几乎相同的流动段，让编码器关注真正有意义的流量变化。
        new_dt, new_q = [], []
        cur_dt, cur_q = float(dt[0]), float(qq[0])

        def close(a: float, b: float) -> bool:
            denom = max(abs(a), abs(b), 1.0)
            return abs(a - b) / denom <= merge_rel_tol

        for i in range(1, dt.size):
            if close(float(qq[i]), cur_q):
                cur_dt += float(dt[i])
            else:
                new_dt.append(cur_dt)
                new_q.append(cur_q)
                cur_dt = float(dt[i])
                cur_q = float(qq[i])
        new_dt.append(cur_dt)
        new_q.append(cur_q)
        dt = np.asarray(new_dt, dtype=np.float64)
        qq = np.asarray(new_q, dtype=np.float64)

    if remove_shutin and dt.size >= 2:
        m = qq > q_thr
        if int(np.sum(m)) >= 2:
            dt = dt[m]
            qq = qq[m]

    return dt, qq, has_shutin


def _make_u_grid(cfg: Config, T_total: float, Nu: int, mode: str) -> np.ndarray:
    t_min = float((cfg.raw["schedule"].get("obs_window", {}) or {}).get("t_min", 1e-6))
    T_total = max(float(T_total), t_min * 2.0)
    if mode == "linear":
        return np.linspace(t_min, T_total, Nu, dtype=np.float64)
    return np.geomspace(t_min, T_total, Nu, dtype=np.float64)


def encode_schedule_to_timegrid(
    cfg: Config,
    sectionIndex: int,
    timeQ: List[float],
    q: List[float],
    n_sections: Optional[int] = None,
) -> EncodedSchedule:
    """将变长流量制度编码为固定长度的神经网络输入。"""
    dt, qq, has_shutin = canonicalize_schedule(cfg, timeQ, q)
    if dt.size < 2 or qq.size != dt.size:
        raise ValueError("invalid schedule after canonicalize")

    sec = int(max(1, sectionIndex))
    N = int(n_sections) if n_sections is not None else int(dt.size)
    N = max(N, int(dt.size))

    T_total = float(np.sum(dt))
    enc = cfg.raw["timegrid_encoding"]
    Nu = int(enc.get("n_u_points", 256))
    grid_mode = str(enc.get("grid", "log")).lower()
    include_cum = bool(enc.get("include_cum", True))
    include_dq = bool(enc.get("include_dq", True))
    include_shutin = bool(enc.get("include_shutin", False))
    q_eps = float(enc.get("q_eps", 1e-12))

    t_edges = np.concatenate([[0.0], np.cumsum(dt)], axis=0)
    u = _make_u_grid(cfg, T_total=T_total, Nu=Nu, mode=("linear" if grid_mode == "linear" else "log"))
    idx = np.searchsorted(t_edges[1:], u, side="right")
    idx = np.clip(idx, 0, dt.size - 1)
    q_u = qq[idx].astype(np.float64)

    # q(t) 是主通道；cum(t)、dq(t) 和关井标志是可选辅助通道，
    # 用于帮助代理模型区分不同流量制度形态。
    cum_u = None
    if include_cum:
        prefix = np.concatenate([[0.0], np.cumsum(dt * qq)], axis=0)
        cum_u = prefix[idx] + (u - t_edges[:-1][idx]) * qq[idx]

    dq_u = None
    if include_dq:
        dq_u = np.zeros_like(u, dtype=np.float64)
        dq_u[1:] = np.diff(q_u)
        dq_u[0] = dq_u[1] if dq_u.size > 1 else 0.0

    shut_u = None
    if include_shutin:
        thr = float(enc.get("shutin_thr", 1e-6))
        shut_u = (q_u <= thr).astype(np.float64)

    norm_mode = str(enc.get("normalize_mode", "global")).lower()
    if norm_mode == "per_sample":
        q_scale = max(float(np.max(q_u)), q_eps)
        cum_scale = max(float(cum_u.max()) if cum_u is not None else 1.0, q_eps)
    else:
        q_scale = float(enc.get("q_global_max", 1.0))
        cum_scale = float(enc.get("cum_global_max", 1.0))

    q_u = q_u / max(q_scale, q_eps)
    if cum_u is not None:
        cum_u = cum_u / max(cum_scale, q_eps)
    if dq_u is not None:
        dq_u = dq_u / max(q_scale, q_eps)

    chans = [q_u]
    if include_cum:
        chans.append(cum_u if cum_u is not None else np.zeros_like(q_u))
    if include_dq:
        chans.append(dq_u if dq_u is not None else np.zeros_like(q_u))
    if include_shutin:
        chans.append(shut_u if shut_u is not None else np.zeros_like(q_u))

    X = np.stack(chans, axis=1)
    x_sched = X.reshape(-1).astype(np.float32)

    # 时间网格之外额外拼接 5 个流动段级别特征，
    # 用于告诉模型当前预测的是哪个流动段对应的双对数曲线。
    n_sec = int(max(1, N))
    sec_clamped = int(np.clip(sec, 1, n_sec))
    section_pos = float((sec_clamped - 1) / max(n_sec - 1, 1))
    sectionIndex_norm = float(sec_clamped / max(n_sec, 1))
    n_sections_norm = float(n_sec / max(12, 1))
    log_T_total = float(np.log(max(T_total, 1e-12)))

    x_sec = np.array(
        [sectionIndex_norm, section_pos, n_sections_norm, float(has_shutin), log_T_total],
        dtype=np.float32,
    )

    return EncodedSchedule(x_sched=x_sched, x_sec=x_sec)