"""Schemi di peso per i portafogli. Ogni funzione ritorna {sleeve_id: peso} con somma 1."""
from __future__ import annotations

import numpy as np
import pandas as pd

_PREFIX = [("PR_", "PAIRS"), ("SH_", "SHAPE"), ("TSM", "TSM"), ("MR", "FADE")]


def family_of(sleeve_id: str) -> str:
    for pre, fam in _PREFIX:
        if sleeve_id.startswith(pre):
            return fam
    return "HONEST"


def _normalize(w: dict[str, float]) -> dict[str, float]:
    tot = sum(w.values())
    return {k: (v / tot if tot > 0 else 0.0) for k, v in w.items()}


def equal(ids: list[str]) -> dict[str, float]:
    n = len(ids)
    return {i: 1.0 / n for i in ids} if n else {}


def manual(ids: list[str], weights: dict[str, float]) -> dict[str, float]:
    return _normalize({i: float(weights.get(i, 0.0)) for i in ids})


def cap(ids: list[str], caps: dict[str, float]) -> dict[str, float]:
    """Equal-weight con tetto al peso AGGREGATO di una famiglia; l'eccesso ridistribuito
    pro-quota alle famiglie non cappate (iterativo finché tutti i cap sono rispettati)."""
    w = equal(ids)
    fam = {i: family_of(i) for i in ids}
    for _ in range(10):
        over = {}
        for f, lim in caps.items():
            members = [i for i in ids if fam[i] == f]
            cur = sum(w[i] for i in members)
            if cur > lim + 1e-12 and members:
                over[f] = (members, lim, cur)
        if not over:
            break
        free_ids = [i for i in ids if fam[i] not in caps]
        freed = 0.0
        for f, (members, lim, cur) in over.items():
            scale = lim / cur
            for i in members:
                freed += w[i] * (1 - scale)
                w[i] *= scale
        if free_ids and freed > 0:
            add = freed / len(free_ids)
            for i in free_ids:
                w[i] += add
        else:
            break
    return _normalize(w)


def inverse_vol(ids: list[str], returns_df: pd.DataFrame, lookback: int) -> dict[str, float]:
    sub = returns_df[ids].iloc[-lookback:]
    vol = sub.std()
    inv = {i: (1.0 / vol[i] if vol[i] and vol[i] > 0 else 0.0) for i in ids}
    return _normalize(inv)


def cluster_rp(ids: list[str], clusters: dict[str, str],
               returns_df: pd.DataFrame, lookback: int) -> dict[str, float]:
    """Equal fra i cluster naturali, poi inverse-vol dentro ogni cluster."""
    groups: dict[str, list[str]] = {}
    for i in ids:
        groups.setdefault(clusters.get(i, i), []).append(i)
    per = 1.0 / len(groups) if groups else 0.0
    w: dict[str, float] = {}
    for members in groups.values():
        iv = inverse_vol(members, returns_df, lookback)
        for i in members:
            w[i] = per * iv[i]
    return _normalize(w)


def weight_vector(scheme: str, ids: list[str], returns_df: pd.DataFrame | None = None,
                  *, weights: dict | None = None, caps: dict | None = None,
                  clusters: dict | None = None, lookback: int = 90) -> dict[str, float]:
    if scheme == "equal":
        return equal(ids)
    if scheme == "manual":
        return manual(ids, weights or {})
    if scheme == "cap":
        return cap(ids, caps or {})
    if scheme == "inverse_vol":
        return inverse_vol(ids, returns_df, lookback)
    if scheme == "cluster_rp":
        return cluster_rp(ids, clusters or {}, returns_df, lookback)
    raise ValueError(f"schema peso sconosciuto: {scheme}")