#!/usr/bin/env python3
# analysis.py

import numpy as np
import pandas as pd
import statsmodels.api as sm
import matplotlib.pyplot as plt

def main():
    # Wczytanie danych
    df = pd.read_json("../data/data.json")

    # Definiujemy X i y do wielozmiennej regresji (możemy się nimi posłużyć w analizie statystycznej)
    X_multi = df[['DL', 'L2_Exposure', 'AoA', 'Usage_Frequency', 'Social_Support', 'Interference']]
    y = df['L2_Competence']
    X_multi = sm.add_constant(X_multi)

    # Wielowymiarowy model (tu raczej do publikacji w tabeli ze współczynnikami)
    model_multi = sm.OLS(y, X_multi).fit()
    print(model_multi.summary())

    # --- [1] WYKRES: L2_Competence vs. L2_Exposure z dopasowaną prostą ---
    # Jednowymiarowa regresja względem L2_Exposure
    X_exp = sm.add_constant(df['L2_Exposure'])
    model_exp = sm.OLS(y, X_exp).fit()
    y_pred_exp = model_exp.predict(X_exp)

    # Sortujemy, by linia była ciągła
    sort_idx = np.argsort(df['L2_Exposure'])
    X_sorted_exp = df['L2_Exposure'].values[sort_idx]
    y_pred_sorted_exp = y_pred_exp.values[sort_idx]

    plt.figure()
    plt.scatter(df['L2_Exposure'], df['L2_Competence'])
    plt.plot(X_sorted_exp, y_pred_sorted_exp)  # Linia regresji
    plt.xlabel('L2 Exposure (h/tyg.)')
    plt.ylabel('L2 Competence (punkty)')
    plt.title('Kompetencja L2 vs. Ekspozycja na język drugi\n(z dopasowaną linią regresji)')
    plt.savefig('../figs/figure_1.png', dpi=300)
    plt.close()

    # --- [2] WYKRES: L2_Competence vs. DL z dopasowaną prostą ---
    # Jednowymiarowa regresja względem DL
    X_dl = sm.add_constant(df['DL'])
    model_dl = sm.OLS(y, X_dl).fit()
    y_pred_dl = model_dl.predict(X_dl)

    # Sortowanie
    sort_idx_dl = np.argsort(df['DL'])
    X_sorted_dl = df['DL'].values[sort_idx_dl]
    y_pred_sorted_dl = y_pred_dl.values[sort_idx_dl]

    plt.figure()
    plt.scatter(df['DL'], df['L2_Competence'])
    plt.plot(X_sorted_dl, y_pred_sorted_dl)  # Linia regresji
    plt.xlabel('Dominujący język (h/tyg.)')
    plt.ylabel('L2 Competence (punkty)')
    plt.title('Kompetencja L2 vs. Czas w języku dominującym\n(z dopasowaną linią regresji)')
    plt.savefig('../figs/figure_2.png', dpi=300)
    plt.close()

    # Eksport tabeli z danymi do LaTeX (opcjonalnie)
    with open("../doc/inc/table.tex", "w", encoding="utf-8") as f:
        f.write(df.to_latex(index=False, 
                            caption="Przykładowe dane dotyczące kompetencji językowych.", 
                            label="tab:data"))
    
    print("\n[INFO] Zapisano wykresy (figure_1.png, figure_2.png) z dopasowanymi liniami regresji.")
    print("[INFO] Zapisano tabelę (table.tex) w folderze figs/.")

if __name__ == "__main__":
    main()