from __future__ import annotations

import hashlib
import json
import math
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any, Mapping, Sequence

import numpy as np

from .features import FEATURE_NAMES, feature_vector
from .io import write_json


def _sigmoid(value: np.ndarray) -> np.ndarray:
    clipped = np.clip(np.asarray(value, dtype=np.float64), -35.0, 35.0)
    return 1.0 / (1.0 + np.exp(-clipped))


def _canonical_bytes(value: Mapping[str, Any]) -> bytes:
    return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":"), allow_nan=False).encode("utf-8")


@dataclass(frozen=True)
class LogisticModel:
    schemaVersion: int
    modelId: str
    modelType: str
    featureNames: list[str]
    classes: list[str]
    mean: list[float]
    scale: list[float]
    weights: list[float]
    intercept: float
    l2: float
    iterations: int
    finalTrainingLoss: float
    lowThreshold: float
    highThreshold: float
    oodZLimit: float
    oodFractionLimit: float
    oodDistanceLimit: float
    trainingDataProvenance: str
    promotionStatus: str
    intendedUse: str
    forbiddenUses: list[str]

    def validate(self) -> None:
        width = len(FEATURE_NAMES)
        if self.schemaVersion != 1 or self.modelType != "l2_logistic_binary_v1":
            raise ValueError("Unsupported model schema or type")
        if tuple(self.featureNames) != FEATURE_NAMES:
            raise ValueError("Model feature contract differs from this package")
        if self.classes != ["no_flag", "review_required"]:
            raise ValueError("Model classes differ from the published decision boundary")
        if any(len(value) != width for value in (self.mean, self.scale, self.weights)):
            raise ValueError("Model parameter dimensions do not match the feature contract")
        arrays = [np.asarray(value, dtype=np.float64) for value in (self.mean, self.scale, self.weights)]
        if not all(np.isfinite(value).all() for value in arrays) or not math.isfinite(self.intercept):
            raise ValueError("Model contains non-finite parameters")
        if any(value <= 0 for value in self.scale):
            raise ValueError("Model scale values must be positive")
        if not 0 <= self.lowThreshold < self.highThreshold <= 1:
            raise ValueError("Model abstention thresholds are invalid")
        if self.promotionStatus not in {"RESEARCH_ONLY", "FIELD_VALIDATION_CANDIDATE"}:
            raise ValueError("Unsupported promotion status")

    def decision_function(self, matrix: np.ndarray) -> np.ndarray:
        self.validate()
        value = np.asarray(matrix, dtype=np.float64)
        if value.ndim == 1:
            value = value[None, :]
        if value.ndim != 2 or value.shape[1] != len(FEATURE_NAMES) or not np.isfinite(value).all():
            raise ValueError("Inference matrix differs from the model feature contract")
        standardised = (value - np.asarray(self.mean)) / np.asarray(self.scale)
        return standardised @ np.asarray(self.weights) + self.intercept

    def predict_proba(self, matrix: np.ndarray) -> np.ndarray:
        positive = _sigmoid(self.decision_function(matrix))
        return np.column_stack([1.0 - positive, positive])

    def ood_metrics(self, matrix: np.ndarray) -> dict[str, Any]:
        value = np.asarray(matrix, dtype=np.float64)
        if value.ndim == 1:
            value = value[None, :]
        z = np.abs((value - np.asarray(self.mean)) / np.asarray(self.scale))
        fraction = (z > self.oodZLimit).mean(axis=1)
        distance = np.sqrt(np.mean(np.square(np.minimum(z, 25.0)), axis=1))
        rejected = (fraction > self.oodFractionLimit) | (distance > self.oodDistanceLimit)
        return {
            "rejected": rejected,
            "fractionBeyondZLimit": fraction,
            "standardisedRmsDistance": distance,
            "zLimit": self.oodZLimit,
            "fractionLimit": self.oodFractionLimit,
            "distanceLimit": self.oodDistanceLimit,
        }


def fit_logistic(
    matrix: np.ndarray,
    labels: np.ndarray,
    *,
    model_id: str = "M3M-DISEASE-SCREEN-SYNTHETIC-V1",
    data_provenance: str = "synthetic",
    promotion_receipt: Mapping[str, Any] | None = None,
    l2: float = 0.05,
    learning_rate: float = 0.12,
    maximum_iterations: int = 3000,
    tolerance: float = 1e-10,
    low_threshold: float = 0.35,
    high_threshold: float = 0.65,
) -> LogisticModel:
    value = np.asarray(matrix, dtype=np.float64)
    target = np.asarray(labels, dtype=np.float64)
    if value.ndim != 2 or value.shape[1] != len(FEATURE_NAMES) or len(target) != len(value):
        raise ValueError("Training data differ from the feature contract")
    if not np.isfinite(value).all() or not np.isfinite(target).all() or set(target.tolist()) != {0.0, 1.0}:
        raise ValueError("Training requires finite features and both binary classes")
    if l2 < 0 or learning_rate <= 0 or maximum_iterations < 1:
        raise ValueError("Invalid optimisation parameters")
    promotion_status = "RESEARCH_ONLY"
    if (
        data_provenance == "field"
        and promotion_receipt is not None
        and promotion_receipt.get("status") == "FIELD_VALIDATION_CANDIDATE"
        and not promotion_receipt.get("failures")
    ):
        promotion_status = "FIELD_VALIDATION_CANDIDATE"
    mean = value.mean(axis=0)
    scale = value.std(axis=0)
    scale = np.where(scale < 1e-6, 1.0, scale)
    x = (value - mean) / scale
    weights = np.zeros(x.shape[1], dtype=np.float64)
    intercept = 0.0
    number_positive = max(1, int((target == 1).sum()))
    number_negative = max(1, int((target == 0).sum()))
    sample_weight = np.where(target == 1, len(target) / (2 * number_positive), len(target) / (2 * number_negative))
    normaliser = float(sample_weight.sum())
    previous_loss = float("inf")
    final_loss = float("inf")
    iterations = 0
    for iteration in range(1, maximum_iterations + 1):
        probability = _sigmoid(x @ weights + intercept)
        error = (probability - target) * sample_weight
        gradient_weights = x.T @ error / normaliser + l2 * weights
        gradient_intercept = float(error.sum() / normaliser)
        weights -= learning_rate * gradient_weights
        intercept -= learning_rate * gradient_intercept
        probability = np.clip(_sigmoid(x @ weights + intercept), 1e-12, 1 - 1e-12)
        final_loss = float(
            -np.sum(sample_weight * (target * np.log(probability) + (1 - target) * np.log(1 - probability))) / normaliser
            + 0.5 * l2 * float(weights @ weights)
        )
        iterations = iteration
        if abs(previous_loss - final_loss) < tolerance:
            break
        previous_loss = final_loss
    if not math.isfinite(final_loss):
        raise RuntimeError("Model optimisation produced a non-finite loss")
    model = LogisticModel(
        schemaVersion=1,
        modelId=model_id,
        modelType="l2_logistic_binary_v1",
        featureNames=list(FEATURE_NAMES),
        classes=["no_flag", "review_required"],
        mean=[round(float(item), 12) for item in mean],
        scale=[round(float(item), 12) for item in scale],
        weights=[round(float(item), 12) for item in weights],
        intercept=round(float(intercept), 12),
        l2=float(l2),
        iterations=iterations,
        finalTrainingLoss=round(final_loss, 12),
        lowThreshold=float(low_threshold),
        highThreshold=float(high_threshold),
        oodZLimit=6.0,
        oodFractionLimit=0.20,
        oodDistanceLimit=3.5,
        trainingDataProvenance=data_provenance,
        promotionStatus=promotion_status,
        intendedUse="Post-flight optical screening to prioritise human ground inspection in a research workflow.",
        forbiddenUses=[
            "Plant disease confirmation",
            "Autonomous flight control",
            "Pesticide, fertiliser or irrigation prescription",
            "Claiming field accuracy from synthetic fixtures",
        ],
    )
    model.validate()
    return model


def model_payload(model: LogisticModel) -> dict[str, Any]:
    model.validate()
    return asdict(model)


def model_sha256(model: LogisticModel) -> str:
    return hashlib.sha256(_canonical_bytes(model_payload(model))).hexdigest()


def save_model(path: str | Path, model: LogisticModel) -> Path:
    payload = model_payload(model)
    payload["modelSha256"] = hashlib.sha256(_canonical_bytes(payload)).hexdigest()
    return write_json(path, payload)


def load_model(path: str | Path) -> LogisticModel:
    payload = json.loads(Path(path).read_text(encoding="utf-8"))
    if not isinstance(payload, dict):
        raise ValueError("Model JSON must be an object")
    expected = payload.pop("modelSha256", None)
    actual = hashlib.sha256(_canonical_bytes(payload)).hexdigest()
    if expected != actual:
        raise ValueError("Model JSON checksum does not match its parameters")
    model = LogisticModel(**payload)
    model.validate()
    return model


def evaluate_binary(model: LogisticModel, matrix: np.ndarray, labels: np.ndarray) -> dict[str, Any]:
    target = np.asarray(labels, dtype=np.int64)
    probability = model.predict_proba(matrix)[:, 1]
    predicted = (probability >= 0.5).astype(np.int64)
    tp = int(((target == 1) & (predicted == 1)).sum())
    tn = int(((target == 0) & (predicted == 0)).sum())
    fp = int(((target == 0) & (predicted == 1)).sum())
    fn = int(((target == 1) & (predicted == 0)).sum())

    def ratio(numerator: float, denominator: float) -> float | None:
        return round(float(numerator / denominator), 6) if denominator else None

    order = np.argsort(-probability, kind="stable")
    ordered_target = target[order]
    positive_count = int((target == 1).sum())
    precision_at_rank = np.cumsum(ordered_target) / np.arange(1, len(target) + 1)
    average_precision = float((precision_at_rank * ordered_target).sum() / positive_count) if positive_count else float("nan")
    return {
        "schemaVersion": 1,
        "sampleCount": len(target),
        "confusion": {"truePositive": tp, "trueNegative": tn, "falsePositive": fp, "falseNegative": fn},
        "metrics": {
            "accuracy": ratio(tp + tn, len(target)),
            "precision": ratio(tp, tp + fp),
            "recall": ratio(tp, tp + fn),
            "specificity": ratio(tn, tn + fp),
            "f1": ratio(2 * tp, 2 * tp + fp + fn),
            "brier": round(float(np.mean(np.square(probability - target))), 6),
            "averagePrecision": round(average_precision, 6) if math.isfinite(average_precision) else None,
        },
        "engineeringSyntheticOnly": model.trainingDataProvenance == "synthetic",
        "warning": "Synthetic metrics verify deterministic software behaviour and must not be reported as field plant-disease accuracy.",
    }


def build_model_card(
    model: LogisticModel,
    *,
    split_receipt: Mapping[str, Any],
    evaluation: Mapping[str, Any],
    promotion: Mapping[str, Any],
) -> dict[str, Any]:
    if promotion.get("status") != model.promotionStatus:
        raise ValueError("Promotion receipt and serialized model status differ")
    return {
        "schemaVersion": 1,
        "modelId": model.modelId,
        "modelSha256": model_sha256(model),
        "modelType": model.modelType,
        "featureContract": model.featureNames,
        "trainingDataProvenance": model.trainingDataProvenance,
        "promotionStatus": promotion["status"],
        "promotionFailures": promotion.get("failures", []),
        "splitReceipt": dict(split_receipt),
        "evaluation": dict(evaluation),
        "thresholds": {
            "noFlagMaximum": model.lowThreshold,
            "reviewRequiredMinimum": model.highThreshold,
            "betweenThresholds": "abstain_and_require_human_review",
        },
        "intendedUse": model.intendedUse,
        "forbiddenUses": model.forbiddenUses,
        "fieldPerformanceEstablished": False,
        "warning": "This model card describes a synthetic engineering prototype, not a clinically or agronomically validated disease diagnostic model.",
    }
