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


def ausgleichung_global(A, dl, Q_ext):
    A=np.asarray(A, float)
    dl = np.asarray(dl, float).reshape(-1, 1)
    Q_ext = np.asarray(Q_ext, float)

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

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

    # 3) Zuschlagsvektor dx und Unbekanntenvektor x
    dx = np.linalg.solve(N, n)

    # 4) Residuenvektor v
    v = A @ dx - dl

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

    # 6) Kofaktormatrix der Beobachtungen Q_ll_dach
    Q_ll_dach = StochastischesModell.berechne_Q_ll_dach(A, Q_xx)

    # 7) Kofaktormatrix der Verbesserungen Q_vv
    Q_vv = StochastischesModell.berechne_Qvv(Q_ext, Q_ll_dach)

    # 8) 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,
    }
    return dict_ausgleichung, dx


def ausgleichung_lokal(
    A,
    dl,
    Q_ll,
    x0,
    liste_punktnummern,
    auswahl,
    mit_massstab: bool = True,
):

    A = np.asarray(A, dtype=float)
    dl = np.asarray(dl, dtype=float).reshape(-1, 1)
    Q_ll = np.asarray(Q_ll, dtype=float)
    x0 = np.asarray(x0, dtype=float).reshape(-1, 1)

    # 1) Gewichtsmatrix
    P = np.linalg.inv(Q_ll)

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

    # 3) Datum: G, E, Gi
    # -> Datumsfestlegung ist sympy-basiert, daher nur dafür kurz Sympy verwenden
    x0_sp = sp.Matrix(x0)
    G = Datumsfestlegung.raenderungsmatrix_G(x0_sp, liste_punktnummern, mit_massstab=mit_massstab)
    aktive = Datumsfestlegung.datumskomponenten(auswahl, liste_punktnummern)
    E = Datumsfestlegung.auswahlmatrix_E(u=A.shape[1], aktive_unbekannte_indices=aktive)
    Gi_sp = E * G

    Gi = np.asarray(Gi_sp, dtype=float)  # zurück nach numpy

    # 4) gerändertes System lösen:
    # [N  Gi] [dx] = [n]
    # [GiT 0] [k ]   [0]
    u = N.shape[0]
    d = Gi.shape[1]

    K = np.block([
        [N,      Gi],
        [Gi.T, np.zeros((d, d))]
    ])
    rhs = np.vstack([n, np.zeros((d, 1))])

    sol = np.linalg.solve(K, rhs)
    dx = sol[:u, :]

    # 5) Residuen
    v = dl - A @ dx

    # 6) Qxx (innere Lösung)
    N_inv = np.linalg.inv(N)
    N_inv_G = N_inv @ Gi
    S = Gi.T @ N_inv_G
    print("rank(Gi) =", np.linalg.matrix_rank(Gi))
    print("Gi shape =", Gi.shape)
    print("rank(S)  =", np.linalg.matrix_rank(S))
    print("S shape  =", S.shape)
    S_inv = np.linalg.inv(S)
    Q_xx = N_inv - N_inv_G @ S_inv @ N_inv_G.T

    # 7) Q_lhat_lhat, Q_vv
    Q_lhat_lhat = A @ Q_xx @ A.T
    Q_vv = np.linalg.inv(P) - Q_lhat_lhat

    # 8) Redundanz
    R = Q_vv @ P
    r_vec = np.diag(R).reshape(-1, 1)
    n_beob = A.shape[0]
    u = A.shape[1]
    d = Gi.shape[1]
    r_gesamt = n_beob - u + d

    # 9) sigma0
    vv = float(v.T @ P @ v)
    sigma0_apost = float(np.sqrt(vv / r_gesamt))

    return {
        "dx": dx,
        "v": v,
        "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": np.asarray(G, dtype=float),
        "Gi": Gi,
    }, dx