from Datumsfestlegung import *
from Stochastisches_Modell import StochastischesModell
import sympy as sp
import Export
import Netzqualität_Genauigkeit


def ausgleichung_global(A, dl, stoch_modell: StochastischesModell):

    #Q_ll, P = stoch_modell.berechne_Qll_P()        #Kofaktormatrix und P-Matrix
    P = sp.eye(A.shape[0])
    N = A.T * P * A                                 #Normalgleichungsmatrix N
    Q_xx = N.inv()                                  #Kofaktormatrix der Unbekannten Qxx
    n = A.T * P * dl                                #Absolutgliedvektor n

    dx = N.LUsolve(n)                               #Zuschlagsvektor dx

    v = dl - A * dx                                 #Residuenvektor v

    Q_ll_dach = A * Q_xx * A.T
    Q_vv = stoch_modell.berechne_Qvv(A, P, Q_xx)          #Kofaktormatrix der Verbesserungen Qvv
    R = stoch_modell.berechne_R(Q_vv, P)                  #Redundanzmatrix R
    r = stoch_modell.berechne_r(R)                        #Redundanzanteile als Vektor r
    redundanzanteile = A.shape[0] - A.shape[1] #n-u+d
    soaposteriori = Netzqualität_Genauigkeit.Genauigkeitsmaße.s0apost(v, P, redundanzanteile)

    dict_ausgleichung = {
        "dx": dx,
        "v": v,
        "P": P,
        "N": N,
        "Q_xx": Q_xx,
        "Q_ll_dach": Q_ll_dach,
        "Q_vv": Q_vv,
        "R": R,
        "r": r,
        "soaposteriori": soaposteriori,
    }

    Export.Export.ausgleichung_to_datei(r"Zwischenergebnisse\Ausgleichung_Iteration0.csv", dict_ausgleichung)

    return dict_ausgleichung, dx



def ausgleichung_lokal(
    A: sp.Matrix,
    dl: sp.Matrix,
    stoch_modell: StochastischesModell,
    x0: sp.Matrix,
    idx_X, idx_Y, idx_Z,
    aktive_unbekannte_indices,
    mit_massstab: bool = True,
):
    # 1) Gewichte
    Q_ll, P = stoch_modell.berechne_Qll_P()
    # Debug-Option:
    # P = sp.eye(A.rows)

    # 2) Normalgleichungen
    N = A.T * P * A
    n = A.T * P * dl

    # 3) Datum (G, E, Gi)
    G = raenderungsmatrix_G(x0, idx_X, idx_Y, idx_Z, mit_massstab=mit_massstab)
    E = auswahlmatrix_E(u=A.cols, aktive_unbekannte_indices=aktive_unbekannte_indices)
    Gi = E * G

    # 4) Geränderte Lösung (dx)
    dx = berechne_dx_geraendert(N, n, Gi)

    # 5) Residuen
    v = dl - A * dx

    # 6) KORREKTE Q_xx für gerändertes Problem:
    # Q_xx = N^{-1} - N^{-1}Gi (Gi^T N^{-1} Gi)^{-1} Gi^T N^{-1}
    # numerisch besser via LUsolve statt inv:
    N_inv = N.inv()  # wenn N groß ist, kann man das unten auch ohne inv machen (siehe Hinweis)
    N_inv_G = N_inv * Gi
    S = Gi.T * N_inv_G
    S_inv = S.inv()
    Q_xx = N_inv - N_inv_G * S_inv * N_inv_G.T

    # 7) Q_lhat_lhat und Q_vv
    Q_lhat_lhat = A * Q_xx * A.T
    Q_vv = P.inv() - Q_lhat_lhat

    # 8) Redundanzmatrix und -anteile
    R = Q_vv * P
    r_vec = sp.Matrix(R.diagonal())

    # 9) Freiheitsgrade (Redundanz gesamt)
    n_beob = A.rows
    u = A.cols
    d = Gi.shape[1]
    r_gesamt = n_beob - u + d

    # 10) sigma0 a posteriori
    omega = float((v.T * P * v)[0, 0])
    sigma0_hat = (omega / float(r_gesamt)) ** 0.5

    return {
        "dx": dx,
        "v": v,
        "Q_ll": Q_ll,
        "P": P,
        "N": N,
        "Q_xx": Q_xx,
        "Q_lhat_lhat": Q_lhat_lhat,
        "Q_vv": Q_vv,
        "R": R,
        "r": r_vec,
        "r_gesamt": r_gesamt,
        "sigma0_hat": sigma0_hat,
        "G": G,
        "Gi": Gi,
    }