from Stochastisches_Modell import StochastischesModell
from Netzqualität_Genauigkeit import Genauigkeitsmaße
from Datumsfestlegung import Datumsfestlegung
import sympy as sp
import Export


def ausgleichung_global(
    A: sp.Matrix,
    dl: sp.Matrix,
    Q_ll: sp.Matrix,
    x0: sp.Matrix,
    idx_X, idx_Y, idx_Z,
    anschluss_indices,
    anschluss_werte,
    Sigma_AA,
):
    # 1) Datumsfestlegung (weiches Datum) System erweitern
    A_ext, dl_ext, Q_ext = Datumsfestlegung.weiches_datum(
        A=A,
        dl=dl,
        Q_ll=Q_ll,
        x0=x0,
        anschluss_indices=anschluss_indices,
        anschluss_werte=anschluss_werte,
        Sigma_AA=Sigma_AA,
    )

    # 2) Gewichtsmatrix P
    P = StochastischesModell.berechne_P(Q_ext)

    # 3) Normalgleichungsmatrix N und Absolutgliedvektor n
    N = A_ext.T * P * A_ext
    n = A_ext.T * P * dl_ext

    # 4) Zuschlagsvektor dx
    dx = N.LUsolve(n)

    # 5) Residuenvektor v
    v = dl - A * dx

    # 6) Kofaktormatrix der Unbekannten Q_xx
    Q_xx = StochastischesModell.berechne_Q_xx(N)

    # 7) Kofaktormatrix der Beobachtungen Q_ll_dach
    Q_ll_dach = A * Q_xx * A.T

    # 8) Kofaktormatrix der Verbesserungen Q_vv
    Q_vv = StochastischesModell.berechne_Qvv(A, P, Q_xx)

    # 9) Redundanzmatrix R und Redundanzanteile r
    R = StochastischesModell.berechne_R(Q_vv, P)                  #Redundanzmatrix R
    r = StochastischesModell.berechne_r(R)                        #Redundanzanteile als Vektor r
    redundanzanteile = A.shape[0] - A.shape[1] #n-u+d

    # 10) s0 a posteriori
    soaposteriori = Genauigkeitsmaße.s0apost(v, P, redundanzanteile)

    # 11) Ausgabe
    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,
    Q_ll: sp.Matrix,
    x0: sp.Matrix,
    idx_X, idx_Y, idx_Z,
    aktive_unbekannte_indices,
    mit_massstab: bool = True,
):
    # 1) Gewichtsmatrix P
    P = StochastischesModell.berechne_P(Q_ll)

    # 2) Normalgleichungsmatrix N und Absolutgliedvektor n
    N = A.T * P * A
    n = A.T * P * dl

    # 3) Datumsfestlegung (Teilspurminimierung)
    G = Datumsfestlegung.raenderungsmatrix_G(x0, liste_punktnummern, mit_massstab=mit_massstab)
    aktive = Datumsfestlegung.datumskomponenten(auswahl, liste_punktnummern)
    E = Datumsfestlegung.auswahlmatrix_E(u=A.cols, aktive_unbekannte_indices=aktive)
    Gi = E * G

    # 4) Zuschlagsvektor dx
    dx = Datumsfestlegung.berechne_dx_geraendert(N, n, Gi)

    # 5) Residuenvektor v
    v = dl - A * dx

    # 6) Kofaktormatrix der Unbekannten Q_xx
    N_inv = N.inv()
    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) Kofaktormatrix der Beobachtungen Q_ll_dach
    Q_lhat_lhat = A * Q_xx * A.T

    # 8) Kofaktormatrix der Verbesserungen Q_vv
    Q_vv = P.inv() - Q_lhat_lhat

    # 9) Redundanzmatrix R, Redundanzanteile r, Redundanz
    R = Q_vv * P
    r_vec = sp.Matrix(R.diagonal())
    n_beob = A.rows
    u = A.cols
    d = Gi.shape[1]
    r_gesamt = n_beob - u + d

    # 10) s0 a posteriori
    sigma0_apost = Genauigkeitsmaße.s0apost(v, P, r_gesamt)

    # 11) Ausgabe
    dict_ausgleichung_lokal = {
        "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_apost": sigma0_apost,
        "G": G,
        "Gi": Gi,
    }

    Export.Export.ausgleichung_to_datei(r"Zwischenergebnisse\Ausgleichung_Iteration0_lokal.csv", dict_ausgleichung_lokal)
    return dict_ausgleichung_lokal, dx