# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Index panel assembly (protocol §17 schema), normalisation, growth rates, reliability grades."""
from __future__ import annotations

import numpy as np
import polars as pl

SCHEMA_COLS = ["geography_id", "geography_name", "geography_level", "property_type", "month", "index_value", "base_period",
               "log_monthly_change", "monthly_pct_change", "annual_pct_change", "actual_transaction_count", "effective_sample_size",
               "local_data_weight", "lower_95", "upper_95", "reliability_grade", "direct_or_model_assisted", "data_cutoff",
               "publication_vintage", "model_version", "index_method", "partial_month_flag", "suppression_flag"]


def normalise(log_index: np.ndarray, months: list[str], base: str = "2021") -> tuple[np.ndarray, str]:
    """2021 average = 100 (geometric mean over the available 2021 months) or a single month = 100."""
    li = np.asarray(log_index, float)
    if len(base) == 4:
        sel = [i for i, m in enumerate(months) if m.startswith(base) and not np.isnan(li[i])]
        if not sel:
            return np.full_like(li, np.nan), f"{base} average = 100 (unavailable)"
        c = float(np.mean(li[sel]))
        label = f"{base} average = 100" + ("" if len(sel) == 12 else f" ({len(sel)} months available)")
    else:
        i = months.index(base); c = li[i]; label = f"{base} = 100"
    return 100.0 * np.exp(li - c), label


def growth(log_index: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    li = np.asarray(log_index, float)
    d1 = np.full_like(li, np.nan); d1[1:] = li[1:] - li[:-1]
    d12 = np.full_like(li, np.nan); d12[12:] = li[12:] - li[:-12]
    return d1, 100 * (np.exp(d1) - 1), 100 * (np.exp(d12) - 1)


def grade(n: np.ndarray, ess: np.ndarray, se_log: np.ndarray) -> np.ndarray:
    """A–E reliability grade per month from effective sample size and posterior se of the log level.
    Thresholds are anchored on the Monte Carlo results (results/simulation) — see paper §7."""
    g = np.full(len(n), "E", dtype=object)
    g[(ess >= 15) & (se_log < 0.05)] = "D"
    g[(ess >= 30) & (se_log < 0.035)] = "C"
    g[(ess >= 75) & (se_log < 0.025)] = "B"
    g[(ess >= 150) & (se_log < 0.015)] = "A"
    return g


def build_rows(geo_id: str, geo_name: str, level: str, cls: str, months: list[str], log_index: np.ndarray, var: np.ndarray,
               n: np.ndarray, ess: np.ndarray, weight: np.ndarray, method: str, data_cutoff: str, vintage: str, model_version: str,
               base: str = "2021", partial_last: bool = True, direct_rule: str = "n>0") -> pl.DataFrame:
    idx, base_label = normalise(log_index, months, base)
    d1, p1, p12 = growth(log_index)
    se = np.sqrt(np.maximum(np.asarray(var, float), 0))
    lo = idx * np.exp(-1.96 * se); hi = idx * np.exp(1.96 * se)
    gr = grade(np.asarray(n, float), np.asarray(ess, float), se)
    direct = np.where(np.asarray(n) > 0, "direct", "model_assisted")
    direct = np.where(np.asarray(ess) < 5, "insufficiently_identified", direct)
    suppress = (np.asarray(ess) < 5) | np.isnan(idx)
    partial = np.zeros(len(months), bool)
    if partial_last:
        partial[-1] = True
    return pl.DataFrame({
        "geography_id": geo_id, "geography_name": geo_name, "geography_level": level, "property_type": cls, "month": months,
        "index_value": idx, "base_period": base_label, "log_monthly_change": d1, "monthly_pct_change": p1, "annual_pct_change": p12,
        "actual_transaction_count": np.asarray(n, float).astype(int), "effective_sample_size": np.asarray(ess, float),
        "local_data_weight": np.asarray(weight, float), "lower_95": lo, "upper_95": hi, "reliability_grade": gr.astype(str),
        "direct_or_model_assisted": direct.astype(str), "data_cutoff": data_cutoff, "publication_vintage": vintage,
        "model_version": model_version, "index_method": method, "partial_month_flag": partial, "suppression_flag": suppress,
    }).select(SCHEMA_COLS)