#!/usr/bin/env python3
"""
Export the cluster classification model to JSON for use in JavaScript.

Reads dimension names directly from bicorder.json so the model always
stays in sync with the current bicorder structure.

When gradients are renamed in bicorder.json, add the old→new mapping to
COLUMN_RENAMES so the training CSV columns are correctly aligned.

Usage:
  python3 scripts/export_model_for_js.py data/synthetic_20251116/readings.csv
  python3 scripts/export_model_for_js.py data/manual_20260101/readings.csv --output bicorder_model.json
"""

import argparse
import json
from pathlib import Path

import pandas as pd
import numpy as np
from sklearn.preprocessing import StandardScaler
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis

# Path to bicorder.json (relative to this script)
BICORDER_JSON = Path(__file__).parent.parent.parent / 'bicorder.json'

# Historical column renames: maps old CSV column names → current bicorder.json names.
# Add an entry here whenever gradient terms are renamed in bicorder.json.
COLUMN_RENAMES = {
    'Design_elite_vs_vernacular': 'Design_institutional_vs_vernacular',
    'Entanglement_exclusive_vs_non-exclusive': 'Entanglement_monopolistic_vs_pluralistic',
    'Experience_sufficient_vs_insufficient': 'Experience_sufficient_vs_limited',
    'Experience_Kafka_vs_Whitehead': 'Experience_restraining_vs_liberating',
}


def load_bicorder_dimensions(bicorder_path):
    """Read DIMENSIONS and KEY_DIMENSIONS from bicorder.json."""
    with open(bicorder_path) as f:
        data = json.load(f)

    dimensions = []
    key_dimensions = []

    for category in data['diagnostic']:
        set_name = category['set_name']
        for gradient in category['gradients']:
            dim_name = f"{set_name}_{gradient['term_left']}_vs_{gradient['term_right']}"
            dimensions.append(dim_name)
            if gradient.get('shortform', False):
                key_dimensions.append(dim_name)

    return dimensions, key_dimensions, data['version']


def main():
    parser = argparse.ArgumentParser(
        description='Export cluster classification model to JSON for JavaScript',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Example usage:
  python3 scripts/export_model_for_js.py data/synthetic_20251116/readings.csv
  python3 scripts/export_model_for_js.py data/manual_20260101/readings.csv --output bicorder_model.json
        """
    )
    parser.add_argument('input_csv', help='Diagnostic readings CSV (e.g. data/synthetic_20251116/readings.csv)')
    parser.add_argument('--output', default='bicorder_model.json',
                        help='Output model JSON path (default: bicorder_model.json)')
    args = parser.parse_args()

    dataset_dir = Path(args.input_csv).parent
    analysis_dir = dataset_dir / 'analysis'

    # Derive dimensions and version from bicorder.json
    DIMENSIONS, KEY_DIMENSIONS, BICORDER_VERSION = load_bicorder_dimensions(BICORDER_JSON)

    print(f"Loaded bicorder.json v{BICORDER_VERSION}")
    print(f"Dimensions: {len(DIMENSIONS)}, key dimensions: {len(KEY_DIMENSIONS)}")

    # Load data
    df = pd.read_csv(args.input_csv)
    clusters = pd.read_csv(analysis_dir / 'data' / 'kmeans_clusters.csv')

    # Rename old column names to match current bicorder.json
    df = df.rename(columns=COLUMN_RENAMES)

    # Remove duplicates
    df = df.drop_duplicates(subset='Descriptor', keep='first')

    # Merge and clean
    merged = df.merge(clusters, on='Descriptor')
    merged_clean = merged.dropna(subset=DIMENSIONS)

    print(f"Training on {len(merged_clean)} protocols")

    # Prepare training data
    X = merged_clean[DIMENSIONS].values
    y = merged_clean['cluster'].values

    # Fit scaler
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)

    # Fit LDA
    lda = LinearDiscriminantAnalysis(n_components=1)
    lda.fit(X_scaled, y)

    # Calculate cluster centroids in scaled space
    cluster_centroids = {}
    for cluster_id in [1, 2]:
        cluster_data = X_scaled[y == cluster_id]
        cluster_centroids[cluster_id] = cluster_data.mean(axis=0).tolist()

    # Calculate cluster means in original space (for reference)
    cluster_means_original = {}
    for cluster_id in [1, 2]:
        cluster_data_original = X[y == cluster_id]
        cluster_means_original[cluster_id] = cluster_data_original.mean(axis=0).tolist()

    # Build model export
    model = {
        'version': '1.0',
        'bicorder_version': BICORDER_VERSION,
        'generated': pd.Timestamp.now().isoformat(),
        'dimensions': DIMENSIONS,
        'key_dimensions': KEY_DIMENSIONS,
        'cluster_names': {
            '1': 'Relational/Cultural',
            '2': 'Institutional/Bureaucratic'
        },
        'cluster_descriptions': {
            '1': 'Community-based, emergent, voluntary, cultural protocols',
            '2': 'Formal, institutional, top-down, bureaucratic protocols'
        },
        'scaler': {
            'mean': scaler.mean_.tolist(),
            'scale': scaler.scale_.tolist()
        },
        'lda': {
            'coefficients': lda.coef_[0].tolist(),
            'intercept': lda.intercept_[0]
        },
        'cluster_centroids_scaled': cluster_centroids,
        'cluster_means_original': cluster_means_original,
        'thresholds': {
            'confidence_low': 0.6,
            'completeness_low': 0.5,
            'boundary_distance_low': 0.5
        },
        'metadata': {
            'total_protocols': len(merged_clean),
            'cluster_1_count': int((y == 1).sum()),
            'cluster_2_count': int((y == 2).sum()),
        }
    }

    # Save to JSON
    with open(args.output, 'w') as f:
        json.dump(model, f, indent=2)

    print(f"\nModel exported to {args.output}")
    print(f"Bicorder version: {BICORDER_VERSION}")
    print(f"Total dimensions: {len(DIMENSIONS)}")
    print(f"Key dimensions (short form):")
    for dim in KEY_DIMENSIONS:
        print(f"  - {dim}")
    print(f"Model size: {len(json.dumps(model))} bytes")


if __name__ == '__main__':
    main()
