from mua_pipeline import FeatureVectorizer, MUA
from mua_pipeline.RobustRegression import RobustRegression
from sklearn.pipeline import Pipeline
from sklearn.linear_model import LinearRegression

# With final regression (prediction)
pipeline = Pipeline([
    ('vectorize', FeatureVectorizer()),
    ('mua', MUA(...)),
    ('regressor', LinearRegression())  # or RobustRegression()
])

# Without final regression (score-based)
pipeline = Pipeline([
    ('vectorize', FeatureVectorizer()),
    ('mua', MUA(...))
])

from mua_pipeline import FeatureVectorizer, MUA, preprocess, plot_results
from mua_pipeline.robust_regression import RobustRegression

# Raw data — may have missing values and non-standard orientation
connectivity_matrices = ...   # e.g., shape (n_regions, n_regions, n_subjects)
behavioral_scores = ...       # e.g., shape (n_subjects,)

# Clean the data: auto-detects format, removes missing subjects, returns standard format
clean_connectivity, clean_behavioral, removed_indices = preprocess(
    connectivity_matrices,
    behavioral_scores,
    missing_strategy='any',   # Remove subjects with zeros, NaNs, or Infs
    verbose=True
)

from sklearn.pipeline import Pipeline
from sklearn.model_selection import cross_val_predict, cross_val_score
from scipy.stats import pearsonr
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
import numpy as np

# Build the CPM pipeline (using robust regression to match Yale MATLAB implementation)
cpm_pipeline = Pipeline([
    ('vectorize', FeatureVectorizer()),
    ('mua', MUA(
        filter_by_sign=True,           # Separate positive/negative networks
        direction='difference',        # mean(pos) - mean(neg), matches original CPM
        selection_method='pvalue',     # p-value thresholding
        selection_threshold=0.05,      # p < 0.05
        weighting_method='binary',     # Binary weights (+1/−1)
        correlation_type='pearson',    # Pearson correlation
        feature_aggregation='mean',    # Mean of selected features
    )),
    ('regressor', RobustRegression())  # Robust regression (matches Yale MATLAB)
    # Or use LinearRegression() for ordinary least squares (matches Finn tutorial)
])

# Cross-validation
cpm_scores = cross_val_score(cpm_pipeline, brain_data, behavior, cv=10)
cpm_predictions = cross_val_predict(cpm_pipeline, brain_data, behavior, cv=10)
print(f"CPM R² (10-fold CV): {cpm_scores.mean():.3f} ± {cpm_scores.std():.3f}")

# Evaluation
cpm_r, cpm_p = pearsonr(behavior, cpm_predictions)
mae = mean_absolute_error(behavior, cpm_predictions)
rmse = np.sqrt(mean_squared_error(behavior, cpm_predictions))
r2 = r2_score(behavior, cpm_predictions)

# Build the PNRS pipeline (score-based, no final regression)
pnrs_pipeline = Pipeline([
    ('vectorize', FeatureVectorizer()),
    ('mua', MUA(
        filter_by_sign=False,              # Single combined score
        selection_method='all',            # Use all features
        weighting_method='regression',     # Regression-derived β weights
        feature_aggregation='sum',         # Sum of weighted features
    ))
    # No regressor — PNRS produces an association score, not a prediction
])

pnrs_scores = pnrs_pipeline.fit_transform(brain_data, behavior)

# Evaluation: correlate scores with behavior
pnrs_r, pnrs_p = pearsonr(behavior, pnrs_scores.flatten())

# PNRS with final regression
pnrs_pred_pipeline = Pipeline([
    ('vectorize', FeatureVectorizer()),
    ('mua', MUA(
        filter_by_sign=False,
        selection_method='all',
        weighting_method='regression',
        feature_aggregation='sum',
    )),
    ('regressor', LinearRegression())  # Final regression step
])

# Cross-validation
pnrs_scores = cross_val_score(pnrs_pred_pipeline, brain_data, behavior, cv=10)
pnrs_predictions = cross_val_predict(pnrs_pred_pipeline, brain_data, behavior, cv=10)
print(f"PNRS R² (10-fold CV): {pnrs_scores.mean():.3f} ± {pnrs_scores.std():.3f}")

# Evaluation
pnrs_r, pnrs_p = pearsonr(behavior, pnrs_predictions)

import numpy as np

# Load pre-computed CSS matrix from a discovery study
CSS_matrix = ...  # regions × regions matrix
# Extract upper triangle to match vectorized FC edges
CSS = CSS_matrix[np.triu_indices(n_nodes, k=1)]

# Build the PCS pipeline
pcs_pipeline = Pipeline([
    ('vectorize', FeatureVectorizer()),
    ('mua', MUA(
        filter_by_sign=False,
        selection_method='all',
        weighting_method='external',       # Use externally provided weights
        external_weights=CSS,              # Pre-computed CSS vector
        feature_aggregation='mean',        # Weighted mean (PCS formula)
    ))
])

pcs_scores = pcs_pipeline.fit_transform(brain_data, behavior)
pcs_results = pcs_scores.flatten()

# Step 1: Derive CSS (β-weights) from discovery sample
css_pipeline = Pipeline([
    ('vectorize', FeatureVectorizer()),
    ('mua', MUA(
        filter_by_sign=False,
        selection_method='all',
        weighting_method='regression',
        feature_aggregation='sum',
    ))
])

css_pipeline.fit(brain_data_discovery, behavior_discovery)
CSS = css_pipeline.named_steps['mua'].edge_weights_

# Save CSS as a symmetric matrix CSV
vectorizer = css_pipeline.named_steps['vectorize']
CSS_matrix = vectorizer.inverse_transform(CSS.reshape(1, -1))[0]
pd.DataFrame(CSS_matrix, index=region_labels, columns=region_labels).to_csv('css_beta_weights.csv')

# Step 2: Apply CSS to compute PCS in validation sample
pcs_pipeline = Pipeline([
    ('vectorize', FeatureVectorizer()),
    ('mua', MUA(
        filter_by_sign=False,
        selection_method='all',
        weighting_method='external',
        external_weights=CSS,
        feature_aggregation='mean',
    ))
])
pcs_scores = pcs_pipeline.fit_transform(brain_data_validation, behavior_validation)

# Step 1: Derive correlations from discovery sample
css_pipeline = Pipeline([
    ('vectorize', FeatureVectorizer()),
    ('mua', MUA(
        filter_by_sign=False,
        selection_method='all',
        weighting_method='correlation',
        feature_aggregation='sum',
    ))
])
css_pipeline.fit(brain_data_discovery, group_labels_discovery)
r = css_pipeline.named_steps['mua'].correlations_

# Convert correlation to Cohen's d
n1 = np.sum(group_labels_discovery == 0)  # e.g., controls
n2 = np.sum(group_labels_discovery == 1)  # e.g., patients
CSS_cohen_d = (r / np.sqrt(1 - r**2)) * np.sqrt(
    ((n1 + n2 - 2) / n1) + ((n1 + n2 - 2) / n2)
)

# Save CSS as a symmetric matrix CSV
vectorizer = css_pipeline.named_steps['vectorize']
CSS_cohen_d_matrix = vectorizer.inverse_transform(CSS_cohen_d.reshape(1, -1))[0]
pd.DataFrame(CSS_cohen_d_matrix, index=region_labels, columns=region_labels).to_csv('css_cohen_d.csv')

# Step 2: Apply CSS to compute PCS in validation sample
pcs_pipeline = Pipeline([
    ('vectorize', FeatureVectorizer()),
    ('mua', MUA(
        filter_by_sign=False,
        selection_method='all',
        weighting_method='external',
        external_weights=CSS_cohen_d,
        feature_aggregation='mean',
    ))
])
pcs_scores = pcs_pipeline.fit_transform(brain_data_validation, group_labels_validation)

plot_results(cpm_predictions, behavioral_scores, title="CPM")
plot_results(pnrs_scores.flatten(), behavioral_scores, title="PNRS")