#!/usr/bin/env python3
"""
Baseline MAE Computation for Wuhan Respiratory Disease Risk Prediction.

Naive baseline: district-level historical mean prediction.
Computes MAE on validation set for 1-day, 3-day, 7-day horizons.
"""

import os
import warnings
warnings.filterwarnings('ignore')

import numpy as np
import pandas as pd
import mlflow
from pathlib import Path

# Paths
PROCESSED_DIR = Path('processed')
OUTPUT_DIR = Path('reports')
OUTPUT_DIR.mkdir(exist_ok=True)

# Train/val split: use first half of available data for train, second half for val
# Medical data starts ~2022-12, so split accordingly
TRAIN_START = '2022-12-01'
TRAIN_END = '2023-06-30'
VAL_START = '2023-07-01'
VAL_END = '2024-12-30'


def load_medical_data():
    """Load and combine outpatient and inpatient data."""
    out = pd.read_csv(PROCESSED_DIR / 'medical' / 'outpatient_daily.csv', parse_dates=['date'])
    inp = pd.read_csv(PROCESSED_DIR / 'medical' / 'inpatient_daily.csv', parse_dates=['date'])

    # Respiratory disease keywords already filtered in ETL
    # Combine: outpatient weight=1, inpatient weight=3 (severity proxy)
    out['weight'] = 1
    inp['weight'] = 3

    combined = pd.concat([
        out[['date', 'district', 'case_count', 'weight']],
        inp[['date', 'district', 'case_count', 'weight']]
    ])

    # Weighted sum per district per day
    combined['weighted_cases'] = combined['case_count'] * combined['weight']
    daily = combined.groupby(['date', 'district']).agg(
        weighted_cases=('weighted_cases', 'sum'),
        case_count=('case_count', 'sum')
    ).reset_index()

    # Normalize: combined score per district per day
    daily['risk_score'] = daily['weighted_cases'] / daily.groupby('district')['weighted_cases'].transform('mean')
    return daily


def load_weather_district_mapping():
    """Load weather station to district mapping from processed weather data."""
    wf = pd.read_parquet(PROCESSED_DIR / 'weather' / 'daily_wuhan_2022.parquet')
    # Map each station to its district
    station_district = wf[['station_id', 'district']].drop_duplicates()
    return station_district


def compute_district_historical_mean(daily, train_start, train_end):
    """Compute historical mean risk score per district for training period."""
    train_data = daily[(daily['date'] >= train_start) & (daily['date'] <= train_end)]
    district_mean = train_data.groupby('district')['risk_score'].mean().reset_index()
    district_mean.columns = ['district', 'predicted_risk']
    return district_mean


def compute_mae(daily, district_predictions, val_start, val_end, horizon_days):
    """
    Compute MAE for a given prediction horizon.

    Args:
        daily: DataFrame with date, district, risk_score
        district_predictions: DataFrame with district, predicted_risk (historical mean)
        val_start, val_end: validation period
        horizon_days: number of days to shift for horizon (0=1-day, 2=3-day, 6=7-day)
    """
    val_data = daily[(daily['date'] >= val_start) & (daily['date'] <= val_end)].copy()
    val_data = val_data.merge(district_predictions, on='district', how='left')
    val_data['predicted_risk'] = val_data['predicted_risk'].fillna(val_data.groupby('district')['risk_score'].transform('mean'))

    # Shift actual values to simulate future prediction
    val_data = val_data.sort_values(['district', 'date'])
    val_data['future_risk'] = val_data.groupby('district')['risk_score'].shift(-horizon_days)
    val_data = val_data.dropna(subset=['future_risk'])

    mae = np.mean(np.abs(val_data['predicted_risk'] - val_data['future_risk']))
    return mae


def main():
    print("Loading medical data...")
    daily = load_medical_data()
    print(f"  Combined daily records: {len(daily)}")
    print(f"  Districts: {daily['district'].nunique()}")
    print(f"  Date range: {daily['date'].min()} to {daily['date'].max()}")

    print(f"\nComputing historical mean baseline...")
    print(f"  Train period: {TRAIN_START} to {TRAIN_END}")
    print(f"  Val period: {VAL_START} to {VAL_END}")

    district_mean = compute_district_historical_mean(daily, TRAIN_START, TRAIN_END)
    print(f"  Districts with baseline: {len(district_mean)}")

    print("\nComputing MAE per horizon...")
    horizons = {'1-day': 0, '3-day': 2, '7-day': 6}
    results = {}
    for name, shift in horizons.items():
        mae = compute_mae(daily, district_mean, VAL_START, VAL_END, shift)
        results[name] = mae
        print(f"  {name} horizon MAE: {mae:.4f}")

    # Save report
    report_path = OUTPUT_DIR / 'baseline_mae.md'
    report = f"""# Baseline MAE Report

## Naive Baseline: District-Level Historical Mean

### Methodology
- **Training period**: {TRAIN_START} to {TRAIN_END}
- **Validation period**: {VAL_START} to {VAL_END}
- **Prediction**: District-level historical mean risk score
- **Risk score**: Weighted combination of outpatient (weight=1) and inpatient (weight=3) case counts, normalized by district mean

### Results

| Horizon | MAE |
|---------|-----|
| 1-day   | {results['1-day']:.4f} |
| 3-day   | {results['3-day']:.4f} |
| 7-day   | {results['7-day']:.4f} |

### Interpretation
- These MAE values represent the error of predicting the historical district mean
- Model must achieve MAE < 0.9x these values to beat the naive baseline
- 1-day horizon should have lowest MAE (most predictable)
- 7-day horizon should have highest MAE (least predictable)
"""
    with open(report_path, 'w') as f:
        f.write(report)
    print(f"\nReport saved to {report_path}")

    # Log to MLflow
    try:
        mlflow.set_experiment("wuhan_respiratory_baseline")
        with mlflow.start_run(run_name="naive_baseline"):
            mlflow.log_param("method", "district_historical_mean")
            mlflow.log_param("train_start", TRAIN_START)
            mlflow.log_param("train_end", TRAIN_END)
            mlflow.log_param("val_start", VAL_START)
            mlflow.log_param("val_end", VAL_END)
            for name, mae in results.items():
                mlflow.log_metric(f"mae_{name.replace('-', '_')}", mae)
            mlflow.log_artifact(report_path)
        print("Logged to MLflow")
    except Exception as e:
        print(f"MLflow logging skipped (server not available): {e}")

    return results


if __name__ == '__main__':
    results = main()
    print("\nDone!")