from __future__ import annotations

import json
import os
from dataclasses import dataclass
from typing import Any, Dict

import numpy as np


@dataclass(frozen=True)
class QuarterV5Metadata:
    raw: Dict[str, Any]
    coeff_min: np.ndarray
    coeff_max: np.ndarray
    coeff_mean: np.ndarray
    coeff_std: np.ndarray
    degree: int
    n_relationships: int
    n_conditions: int


def load_metadata_v5(data_dir: str) -> QuarterV5Metadata:
    path = os.path.join(data_dir, "metadata.json")
    with open(path, "r") as f:
        raw = json.load(f)

    coeff_min = np.asarray(raw["coefficient_min"], dtype=np.float32)
    coeff_max = np.asarray(raw["coefficient_max"], dtype=np.float32)
    coeff_mean = np.asarray(raw["coefficient_mean"], dtype=np.float32)
    coeff_std = np.asarray(raw["coefficient_std"], dtype=np.float32)
    degree = int(raw.get("polynomial_degree", 3))
    n_relationships = 7
    n_conditions = int(n_relationships * degree)
    return QuarterV5Metadata(
        raw=raw,
        coeff_min=coeff_min,
        coeff_max=coeff_max,
        coeff_mean=coeff_mean,
        coeff_std=coeff_std,
        degree=degree,
        n_relationships=n_relationships,
        n_conditions=n_conditions,
    )


def normalize_condition_v5(cond_raw_7x3: np.ndarray, meta: QuarterV5Metadata, method: str) -> np.ndarray:
    x = np.asarray(cond_raw_7x3, dtype=np.float32).reshape(-1)
    method = (method or "zscore").lower()
    if method == "minmax":
        rng = meta.coeff_max - meta.coeff_min
        rng = np.where(rng == 0, 1.0, rng)
        return ((x - meta.coeff_min) / rng).astype(np.float32)
    if method == "zscore":
        std = np.where(meta.coeff_std == 0, 1.0, meta.coeff_std)
        return ((x - meta.coeff_mean) / std).astype(np.float32)
    raise ValueError(f"Unknown normalization method: {method}")