import csv
import json

import numpy as np

from modeling import preprocess

N = 96


def _coords():
    yy, xx = np.meshgrid(np.linspace(0, 1, N), np.linspace(0, 1, N), indexing="ij")
    return xx, yy


def _to_field(v):
    a = np.asarray(v, dtype=np.float64)
    if a.ndim == 0:
        return np.full((N, N), float(a))
    if a.shape == (N, N):
        return a
    yi = np.linspace(0, a.shape[0] - 1, N).round().astype(int)
    xi = np.linspace(0, a.shape[1] - 1, N).round().astype(int)
    return a[yi][:, xi]


def cylinder_mask(rows, cols, radius_frac=0.4):
    xx, yy = _coords()
    m = np.zeros((N, N))
    px, py = 1.0 / cols, 1.0 / rows
    r = radius_frac * min(px, py)
    for i in range(rows):
        cy = (i + 0.5) * py
        for j in range(cols):
            cx = (j + 0.5) * px
            m[(xx - cx) ** 2 + (yy - cy) ** 2 <= r * r] = 1.0
    return m


def rect_mask(aspect=1.0, fill=0.7):
    xx, yy = _coords()
    m = np.zeros((N, N))
    hh = min(0.98, (fill / aspect) ** 0.5)
    ww = min(0.98, aspect * hh)
    x0, x1 = 0.5 - ww / 2, 0.5 + ww / 2
    y0, y1 = 0.5 - hh / 2, 0.5 + hh / 2
    m[(xx >= x0) & (xx <= x1) & (yy >= y0) & (yy <= y1)] = 1.0
    return m


def from_fields(
    mask, q_source_W_m3, k_field_W_mK, h_conv_W_m2K, T_amb_degC, domain_L_m
):
    m = (_to_field(mask) > 0.5).astype(np.float64)
    q = _to_field(q_source_W_m3) * m
    k = _to_field(k_field_W_mK)
    h = _to_field(h_conv_W_m2K)
    ta = _to_field(T_amb_degC)
    xx, yy = _coords()
    logL = np.full((N, N), float(np.log(domain_L_m)))
    field = np.stack([m, q, k, h, ta, xx, yy, logL], 0).astype(np.float32)
    return preprocess(field)


def from_pack(
    mask,
    current_A,
    soc,
    R0_ohm,
    k_cell_W_mK,
    k_coolant_W_mK,
    h_conv_W_m2K,
    T_amb_degC,
    domain_L_m,
    beta=2.0,
):
    m = (_to_field(mask) > 0.5).astype(np.float64)
    hg = domain_L_m / (N - 1)
    R_int = R0_ohm * (1.0 + beta * (1.0 - soc) ** 2)
    P_total = current_A**2 * R_int
    area = max(m.sum() * hg * hg, hg * hg)
    q = m * (P_total / area)
    k = np.where(m > 0, k_cell_W_mK, k_coolant_W_mK)
    return from_fields(m, q, k, h_conv_W_m2K, T_amb_degC, domain_L_m)


def from_params(d):
    if d.get("rows") not in (None, "") and d.get("cols") not in (None, ""):
        mask = cylinder_mask(
            int(float(d["rows"])),
            int(float(d["cols"])),
            float(d.get("radius_frac") or 0.4),
        )
    else:
        mask = rect_mask(float(d.get("aspect") or 1.0), float(d.get("fill") or 0.7))
    return from_pack(
        mask,
        float(d["current_A"]),
        float(d["soc"]),
        float(d["R0_ohm"]),
        float(d["k_cell_W_mK"]),
        float(d["k_coolant_W_mK"]),
        float(d["h_conv_W_m2K"]),
        float(d["T_amb_degC"]),
        float(d["domain_L_m"]),
        beta=float(d.get("beta") or 2.0),
    )


def from_json(path):
    with open(path) as f:
        return from_params(json.load(f))


def from_csv(path):
    with open(path) as f:
        return from_params(next(csv.DictReader(f)))