import sympy as sp
from dataclasses import dataclass, field
from typing import Dict, Tuple, Iterable

@dataclass
class StochastischesModell:
    sigma_beob: Iterable[float]           #σ der einzelnen Beobachtung
    group_beob: Iterable[int]             #Gruppenzugehörigkeit jeder Beobachtung (Distanz, Richtung, GNSS, Nivellement,...,)
    sigma0_groups: Dict[int, float] = field(default_factory=dict)  #σ0² für jede Gruppe


    def __post_init__(self):
        self.sigma_beob = sp.Matrix(list(self.sigma_beob))  #Spaltenvektor
        self.group_beob = sp.Matrix(list(self.group_beob))  #Spaltenvektor

        if self.sigma_beob.rows != self.group_beob.rows:
            raise ValueError("sigma_obs und group_ids müssen gleich viele Einträge haben.")

        unique_groups = sorted({int(g) for g in self.group_beob})       #jede Beobachtungsgruppe wird genau einmal berücksichtigt
        for g in unique_groups:
            if g not in self.sigma0_groups:                             #Fehlende Gruppen mit σ_0j^2 = 1.0
                self.sigma0_groups[g] = 1.0


    @property
    def n_beob(self) -> int:
        return int(self.sigma_beob.rows)


    def berechne_Qll_P(self) -> Tuple[sp.Matrix, sp.Matrix]:
        n = self.n_beob
        Q_ll = sp.zeros(n, n)
        P = sp.zeros(n, n)

        for i in range(self.n_beob):
            sigma_i = self.sigma_beob[i, 0]         #σ-Wert der i-ten Beobachtung holen
            g = int(self.group_beob[i, 0])          #Gruppenzugehörigkeit der Beobachtung bestimmen
            sigma0_sq = self.sigma0_groups[g]       #Den Varianzfaktor der Gruppe holen
            q_ii = sigma_i**2                       #σ² berechnen
            Q_ll[i, i] = q_ii                       #Diagonale
            P[i, i] = 1 / (sigma0_sq * q_ii)        #durch VKS nicht mehr P=Qll^-1
        return Q_ll, P


    def berechne_Qvv(self, A: sp.Matrix, P: sp.Matrix, Q_xx: sp.Matrix) -> sp.Matrix:
        Q_vv = P.inv() - A * Q_xx * A.T
        return Q_vv                                 #Kofaktormatrix der Beobachtungsresiduen


    def berechne_R(self, Q_vv: sp.Matrix, P: sp.Matrix) -> sp.Matrix:
        R = Q_vv * P
        return R                                    #Redundanzmatrix


    def berechne_r(self, R: sp.Matrix) -> sp.Matrix:
        n = R.rows
        r = sp.zeros(n, 1)
        for i in range(n):
            r[i, 0] = R[i, i]
        return r                                    #Redundanzanteile


    def berechne_vks(self,v: sp.Matrix, P: sp.Matrix, r: sp.Matrix) -> Dict[int, float]:
        if v.rows != self.n_beob:
            raise ValueError("v passt nicht zur Anzahl der Beobachtungen.")
        gruppen = sorted({int(g) for g in self.group_beob})
        sigma_gruppen: Dict[int, float] = {}
        for g in gruppen:
            idx = [i for i in range(self.n_beob)
                   if int(self.group_beob[i, 0]) == g]
            if not idx:
                continue

            v_g = sp.Matrix([v[i, 0] for i in idx])
            P_g = sp.zeros(len(idx), len(idx))
            for k, i_beob in enumerate(idx):
                P_g[k, k] = P[i_beob, i_beob]
            r_g = sum(r[i_beob, 0] for i_beob in idx)
            sigma_gruppe_g = (v_g.T * P_g * v_g)[0, 0] / r_g
            sigma_gruppen[g] = float(sigma_gruppe_g)
        return sigma_gruppen


    def update_sigma0_von_vks(self, sigma_hat: Dict[int, float]) -> None:
        for g, val in sigma_hat.items():
            self.sigma0_groups[int(g)] = float(val)