import numpy as np
import pandas as pd
from sklearn.linear_model import Ridge
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import Pipeline
from sklearn.model_selection import train_test_split

# Keys for GRA basis
GRA_KEYS = ['const','phi','phi_inv','fib','dyad','prime_log','omega']
PARAM_KEYS = ['alpha','beta','eta','zeta']

def load_dataset(csv_path: str):
    """
    CSV format: const,phi,phi_inv,fib,dyad,prime_log,omega,alpha,beta,eta,zeta
    """
    df = pd.read_csv(csv_path)
    X = df[GRA_KEYS].values
    Y = df[PARAM_KEYS].values
    return X, Y

def learn_mapping(X, Y, alpha=1e-2):
    """
    Learn T (4x7 matrix) with ridge regression.
    """
    model = Pipeline([
        ('scaler', StandardScaler()),
        ('ridge', Ridge(alpha=alpha, fit_intercept=True))
    ])
    model.fit(X, Y)
    # Extract coefficients and intercept
    ridge = model.named_steps['ridge']
    scaler = model.named_steps['scaler']
    coefs = ridge.coef_ @ np.diag(1/scaler.scale_)
    intercept = ridge.intercept_ - coefs @ scaler.mean_
    return coefs, intercept, model

def predict_params(model, raw_gra: dict):
    """
    raw_gra: dict with GRA components
    returns: (alpha, beta, eta, zeta)
    """
    x = np.array([[raw_gra.get(k, 0.0) for k in GRA_KEYS]])
    y = model.predict(x)
    return dict(zip(PARAM_KEYS, y[0]))

# Example usage
if __name__ == "__main__":
    # Example: generate synthetic dataset (replace with real CSV)
    np.random.seed(42)
    N = 200
    X_fake = np.random.rand(N, len(GRA_KEYS)) * [1,1.6,0.6,21,32,5,0.1]
    true_T = np.array([
        [1.0, 0.9, 0.6, 0.4, 0.05, -0.2, -0.5],
        [0.0, 0.4, 0.9, 0.8, 0.15, -0.1, -0.3],
        [0.1, 0.8, 0.5, 0.2, 0.02, 0.3, 0.4],
        [0.2, 0.3, 0.7, 0.1, 0.20, 0.5, 0.2]
    ])
    Y_fake = X_fake @ true_T.T + np.random.randn(N,4)*0.05

    # Split dataset
    X_train, X_test, Y_train, Y_test = train_test_split(X_fake, Y_fake, test_size=0.2)

    # Learn mapping
    coefs, intercept, model = learn_mapping(X_train, Y_train, alpha=1e-2)
    print("Learned T matrix:\n", coefs)
    print("Intercept:\n", intercept)

    # Evaluate
    score = model.score(X_test, Y_test)
    print("R^2 test score:", score)

    # Predict for a sample raw GRA state
    raw = {'const':1.0,'phi':1.618,'phi_inv':0.618,'fib':21,'dyad':32,'prime_log':np.log(97),'omega':0.05}
    params = predict_params(model, raw)
    print("Predicted cymatic params:", params)
