#!/usr/bin/env python3
"""
Step 2: Data Preprocessing and Quantitative Signal Analysis
============================================================

This script performs:
1. Data loading and cleaning from FAERS raw data
2. Standardization of adverse event terms into three categories
3. Temporal trend analysis by quarter/year
4. Severity metrics calculation
5. Visualizations: time series, severity analysis, event comparison
6. Statistical summary generation

Author: K-Dense Coding Agent
Date: 2026-02-03
"""

import json
import os
from datetime import datetime
from collections import defaultdict
from pathlib import Path

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.dates as mdates

# Set random seed for reproducibility
np.random.seed(42)

# Configure matplotlib for non-interactive mode
plt.switch_backend('Agg')
plt.rcParams['font.family'] = 'sans-serif'
plt.rcParams['font.size'] = 10
plt.rcParams['axes.linewidth'] = 0.8
plt.rcParams['figure.dpi'] = 150

# Define paths
SESSION_DIR = Path('/app/sandbox/session_20260203_091322_5981a70f834a')
DATA_DIR = SESSION_DIR / 'workflow' / 'data'
FIGURES_DIR = SESSION_DIR / 'figures'
RESULTS_DIR = SESSION_DIR / 'results'

# Ensure output directories exist
FIGURES_DIR.mkdir(parents=True, exist_ok=True)
RESULTS_DIR.mkdir(parents=True, exist_ok=True)


def load_faers_data(filepath: str) -> dict:
    """Load FAERS raw data from JSON file."""
    print(f"Loading FAERS data from {filepath}...")
    with open(filepath, 'r') as f:
        data = json.load(f)
    print(f"  - Total records: {data['metadata']['total_records']}")
    print(f"  - Drugs: {data['metadata']['drugs']}")
    return data


def parse_receivedate(date_str: str) -> datetime:
    """Parse FAERS date string (YYYYMMDD) to datetime object."""
    if date_str and len(date_str) >= 8:
        try:
            return datetime.strptime(date_str[:8], '%Y%m%d')
        except ValueError:
            return None
    return None


def get_quarter(dt: datetime) -> str:
    """Get quarter string (e.g., '2020-Q1') from datetime."""
    if dt:
        quarter = (dt.month - 1) // 3 + 1
        return f"{dt.year}-Q{quarter}"
    return None


def categorize_reaction(reaction_term: str, adverse_event_categories: dict) -> str:
    """
    Categorize a reaction term into one of the three primary categories:
    - thyroid_cancer
    - pancreatitis
    - gastroparesis
    """
    reaction_lower = reaction_term.lower()

    # Check thyroid cancer terms
    for term in adverse_event_categories.get('thyroid_cancer', []):
        if term.lower() in reaction_lower or reaction_lower in term.lower():
            return 'Thyroid Cancer'

    # Check pancreatitis terms - be inclusive of related terms
    pancreatitis_terms = adverse_event_categories.get('pancreatitis', [])
    if any(term.lower() in reaction_lower for term in pancreatitis_terms):
        return 'Pancreatitis'
    if 'pancreatitis' in reaction_lower:
        return 'Pancreatitis'

    # Check gastroparesis terms
    for term in adverse_event_categories.get('gastroparesis', []):
        if term.lower() in reaction_lower or reaction_lower in term.lower():
            return 'Gastroparesis'

    return None


def extract_severity(record: dict) -> dict:
    """Extract severity information from a FAERS record."""
    serious = record.get('serious') == '1'
    death = record.get('seriousnessdeath') == '1'
    hospitalization = record.get('seriousnesshospitalization') == '1'

    return {
        'is_serious': serious,
        'is_death': death,
        'is_hospitalization': hospitalization,
        'severity_category': 'Death' if death else ('Hospitalization' if hospitalization else ('Other Serious' if serious else 'Non-Serious'))
    }


def process_records(data: dict) -> pd.DataFrame:
    """Process FAERS records into a cleaned DataFrame."""
    print("\nProcessing records...")
    records_list = []
    adverse_event_categories = data['metadata']['adverse_events']

    total_records = len(data['records'])
    for i, record in enumerate(data['records']):
        if i % 100 == 0:
            print(f"  Processing record {i}/{total_records}...")

        # Parse date
        receive_date = parse_receivedate(record.get('receivedate'))
        quarter = get_quarter(receive_date)
        year = receive_date.year if receive_date else None

        # Get drug name
        drug = record.get('drug', '').capitalize()

        # Get matched categories from Step 1
        matched_categories = record.get('matched_categories', [])

        # Extract severity
        severity = extract_severity(record)

        # Also re-categorize reactions to ensure consistency
        primary_categories = set()
        for reaction in record.get('patient_reactions', []):
            term = reaction.get('reactionmeddrapt', '')
            category = categorize_reaction(term, adverse_event_categories)
            if category:
                primary_categories.add(category)

        # Use matched_categories from Step 1 as primary source
        for cat in matched_categories:
            if cat == 'thyroid_cancer':
                primary_categories.add('Thyroid Cancer')
            elif cat == 'pancreatitis':
                primary_categories.add('Pancreatitis')
            elif cat == 'gastroparesis':
                primary_categories.add('Gastroparesis')

        # Create a record for each matched category
        for category in primary_categories:
            records_list.append({
                'drug': drug,
                'date': receive_date,
                'quarter': quarter,
                'year': year,
                'category': category,
                'is_serious': severity['is_serious'],
                'is_death': severity['is_death'],
                'is_hospitalization': severity['is_hospitalization'],
                'severity_category': severity['severity_category']
            })

    df = pd.DataFrame(records_list)
    print(f"  - Processed {len(df)} drug-category pairs from {total_records} records")
    return df


def calculate_temporal_trends(df: pd.DataFrame) -> pd.DataFrame:
    """Calculate temporal trends by quarter for each drug-event pair."""
    print("\nCalculating temporal trends...")

    # Group by quarter, drug, and category
    trends = df.groupby(['quarter', 'drug', 'category']).size().reset_index(name='count')

    # Add year extraction for sorting
    trends['year'] = trends['quarter'].str.extract(r'(\d{4})').astype(int)
    trends['q_num'] = trends['quarter'].str.extract(r'Q(\d)').astype(int)
    trends['sort_key'] = trends['year'] * 10 + trends['q_num']
    trends = trends.sort_values('sort_key')

    print(f"  - Generated {len(trends)} quarterly data points")
    return trends


def calculate_severity_metrics(df: pd.DataFrame, summary_counts: dict) -> pd.DataFrame:
    """Calculate severity metrics for each drug-event combination."""
    print("\nCalculating severity metrics...")

    # Group by drug, category, and severity
    severity_stats = df.groupby(['drug', 'category', 'severity_category']).size().unstack(fill_value=0)

    # Calculate totals and proportions
    severity_stats['Total'] = severity_stats.sum(axis=1)

    # Calculate proportions
    for col in ['Death', 'Hospitalization', 'Other Serious', 'Non-Serious']:
        if col in severity_stats.columns:
            severity_stats[f'{col}_pct'] = (severity_stats[col] / severity_stats['Total'] * 100).round(2)

    severity_stats = severity_stats.reset_index()
    print(f"  - Calculated severity metrics for {len(severity_stats)} drug-category pairs")
    return severity_stats


def calculate_event_counts(df: pd.DataFrame) -> pd.DataFrame:
    """Calculate raw event counts by drug and category."""
    print("\nCalculating event counts...")

    counts = df.groupby(['drug', 'category']).size().reset_index(name='count')
    counts_pivot = counts.pivot(index='category', columns='drug', values='count').fillna(0).astype(int)

    print(f"  - Event counts calculated")
    return counts_pivot


def calculate_year_over_year_change(df: pd.DataFrame) -> dict:
    """Calculate year-over-year changes in reporting."""
    print("\nCalculating year-over-year changes...")

    # Get yearly counts
    yearly = df.groupby(['year', 'drug', 'category']).size().reset_index(name='count')

    changes = {}
    for drug in df['drug'].unique():
        changes[drug] = {}
        for category in df['category'].unique():
            subset = yearly[(yearly['drug'] == drug) & (yearly['category'] == category)]
            if len(subset) >= 2:
                subset = subset.sort_values('year')
                years = subset['year'].values
                counts = subset['count'].values

                # Calculate percent change between last two complete years
                if len(years) >= 2:
                    latest_year = years[-1]
                    prev_year = years[-2]
                    latest_count = counts[-1]
                    prev_count = counts[-2]

                    if prev_count > 0:
                        pct_change = ((latest_count - prev_count) / prev_count) * 100
                        changes[drug][category] = {
                            'prev_year': int(prev_year),
                            'latest_year': int(latest_year),
                            'prev_count': int(prev_count),
                            'latest_count': int(latest_count),
                            'pct_change': round(pct_change, 1)
                        }

    return changes


def plot_time_series(trends: pd.DataFrame, output_path: str):
    """Create time series plot showing adverse event trends over time."""
    print(f"\nGenerating time series plot: {output_path}")

    fig, axes = plt.subplots(1, 3, figsize=(15, 5))

    drugs = ['Semaglutide', 'Tirzepatide', 'Liraglutide']
    colors = {'Thyroid Cancer': '#e74c3c', 'Pancreatitis': '#3498db', 'Gastroparesis': '#2ecc71'}

    for idx, drug in enumerate(drugs):
        ax = axes[idx]
        drug_data = trends[trends['drug'] == drug]

        for category in ['Thyroid Cancer', 'Pancreatitis', 'Gastroparesis']:
            cat_data = drug_data[drug_data['category'] == category].copy()
            if not cat_data.empty:
                cat_data = cat_data.sort_values('sort_key')
                ax.plot(range(len(cat_data)), cat_data['count'].values,
                       marker='o', markersize=4, label=category,
                       color=colors.get(category, '#333'), linewidth=1.5)

        ax.set_title(f'{drug}', fontsize=12, fontweight='bold')
        ax.set_xlabel('Quarter (chronological order)', fontsize=10)
        ax.set_ylabel('Report Count', fontsize=10)
        ax.legend(loc='upper left', fontsize=8)
        ax.grid(True, alpha=0.3)
        ax.set_ylim(bottom=0)

        # Set x-tick labels to show quarters
        if len(drug_data) > 0:
            quarters = drug_data.drop_duplicates('quarter').sort_values('sort_key')['quarter'].tolist()
            # Show every 4th label to avoid crowding
            tick_positions = range(0, len(quarters), max(1, len(quarters)//6))
            ax.set_xticks(list(tick_positions))
            ax.set_xticklabels([quarters[i] if i < len(quarters) else '' for i in tick_positions],
                              rotation=45, ha='right', fontsize=8)

    plt.suptitle('Adverse Event Reports Over Time by Drug', fontsize=14, fontweight='bold', y=1.02)
    plt.tight_layout()
    plt.savefig(output_path, dpi=300, bbox_inches='tight', facecolor='white')
    plt.close()
    print(f"  - Saved time series plot")


def plot_severity_analysis(df: pd.DataFrame, output_path: str):
    """Create stacked bar chart comparing severity outcomes across drugs."""
    print(f"\nGenerating severity analysis plot: {output_path}")

    # Calculate severity proportions by drug
    severity_by_drug = df.groupby(['drug', 'severity_category']).size().unstack(fill_value=0)
    severity_by_drug = severity_by_drug.div(severity_by_drug.sum(axis=1), axis=0) * 100

    # Ensure consistent column order
    severity_order = ['Death', 'Hospitalization', 'Other Serious', 'Non-Serious']
    for col in severity_order:
        if col not in severity_by_drug.columns:
            severity_by_drug[col] = 0
    severity_by_drug = severity_by_drug[severity_order]

    fig, ax = plt.subplots(figsize=(10, 6))

    colors = ['#c0392b', '#e74c3c', '#f39c12', '#27ae60']

    x = np.arange(len(severity_by_drug.index))
    width = 0.6

    bottom = np.zeros(len(x))
    for i, (col, color) in enumerate(zip(severity_order, colors)):
        values = severity_by_drug[col].values
        ax.bar(x, values, width, label=col, bottom=bottom, color=color)
        bottom += values

    ax.set_xlabel('Drug', fontsize=12)
    ax.set_ylabel('Proportion (%)', fontsize=12)
    ax.set_title('Severity Distribution of Adverse Event Reports by Drug', fontsize=14, fontweight='bold')
    ax.set_xticks(x)
    ax.set_xticklabels(severity_by_drug.index, fontsize=11)
    ax.legend(title='Outcome', bbox_to_anchor=(1.02, 1), loc='upper left')
    ax.set_ylim(0, 100)
    ax.grid(True, axis='y', alpha=0.3)

    # Add percentage labels
    for i, drug in enumerate(severity_by_drug.index):
        cumsum = 0
        for col, color in zip(severity_order, colors):
            val = severity_by_drug.loc[drug, col]
            if val > 5:  # Only label if segment is large enough
                ax.text(i, cumsum + val/2, f'{val:.1f}%', ha='center', va='center',
                       fontsize=9, color='white', fontweight='bold')
            cumsum += val

    plt.tight_layout()
    plt.savefig(output_path, dpi=300, bbox_inches='tight', facecolor='white')
    plt.close()
    print(f"  - Saved severity analysis plot")


def plot_event_comparison(counts_pivot: pd.DataFrame, output_path: str):
    """Create grouped bar chart comparing event counts across drugs."""
    print(f"\nGenerating event comparison plot: {output_path}")

    fig, ax = plt.subplots(figsize=(10, 6))

    categories = counts_pivot.index.tolist()
    drugs = counts_pivot.columns.tolist()

    x = np.arange(len(categories))
    width = 0.25

    colors = ['#3498db', '#e74c3c', '#2ecc71']

    for i, (drug, color) in enumerate(zip(drugs, colors)):
        values = counts_pivot[drug].values
        bars = ax.bar(x + i*width, values, width, label=drug, color=color)

        # Add count labels on bars
        for bar, val in zip(bars, values):
            if val > 0:
                ax.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 2,
                       str(int(val)), ha='center', va='bottom', fontsize=9, fontweight='bold')

    ax.set_xlabel('Adverse Event Category', fontsize=12)
    ax.set_ylabel('Report Count', fontsize=12)
    ax.set_title('Adverse Event Reports by Drug and Category', fontsize=14, fontweight='bold')
    ax.set_xticks(x + width)
    ax.set_xticklabels(categories, fontsize=11)
    ax.legend(title='Drug', loc='upper right')
    ax.grid(True, axis='y', alpha=0.3)

    # Set y-limit with some padding
    max_val = counts_pivot.max().max()
    ax.set_ylim(0, max_val * 1.15)

    plt.tight_layout()
    plt.savefig(output_path, dpi=300, bbox_inches='tight', facecolor='white')
    plt.close()
    print(f"  - Saved event comparison plot")


def generate_summary(df: pd.DataFrame, summary_counts: dict, yoy_changes: dict,
                     severity_stats: pd.DataFrame, output_path: str):
    """Generate a text summary of key statistical findings."""
    print(f"\nGenerating statistical summary: {output_path}")

    lines = []
    lines.append("=" * 70)
    lines.append("GLP-1 RECEPTOR AGONISTS: ADVERSE EVENT ANALYSIS SUMMARY")
    lines.append("=" * 70)
    lines.append(f"\nAnalysis Date: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    lines.append(f"Total Records Analyzed: {len(df)}")

    # Overview by drug
    lines.append("\n" + "-" * 50)
    lines.append("REPORT COUNTS BY DRUG AND ADVERSE EVENT")
    lines.append("-" * 50)

    drug_totals = df.groupby('drug').size()
    for drug in ['Semaglutide', 'Tirzepatide', 'Liraglutide']:
        total = drug_totals.get(drug, 0)
        lines.append(f"\n{drug}: {total} total reports")

        for cat in ['Thyroid Cancer', 'Pancreatitis', 'Gastroparesis']:
            count = len(df[(df['drug'] == drug) & (df['category'] == cat)])
            if count > 0:
                lines.append(f"  - {cat}: {count} reports")

    # Year-over-year trends
    lines.append("\n" + "-" * 50)
    lines.append("YEAR-OVER-YEAR TRENDS")
    lines.append("-" * 50)

    for drug in ['Semaglutide', 'Tirzepatide', 'Liraglutide']:
        if drug in yoy_changes and yoy_changes[drug]:
            lines.append(f"\n{drug}:")
            for cat, data in yoy_changes[drug].items():
                direction = "increase" if data['pct_change'] > 0 else "decrease"
                lines.append(f"  - {cat}: {abs(data['pct_change']):.1f}% {direction} "
                           f"({data['prev_year']}: {data['prev_count']} -> "
                           f"{data['latest_year']}: {data['latest_count']})")

    # Severity analysis
    lines.append("\n" + "-" * 50)
    lines.append("SEVERITY ANALYSIS")
    lines.append("-" * 50)

    severity_by_drug = df.groupby(['drug', 'severity_category']).size().unstack(fill_value=0)
    for drug in ['Semaglutide', 'Tirzepatide', 'Liraglutide']:
        if drug in severity_by_drug.index:
            row = severity_by_drug.loc[drug]
            total = row.sum()
            lines.append(f"\n{drug} (n={total}):")

            for cat in ['Death', 'Hospitalization', 'Other Serious', 'Non-Serious']:
                if cat in row:
                    count = row[cat]
                    pct = (count / total * 100) if total > 0 else 0
                    lines.append(f"  - {cat}: {count} ({pct:.1f}%)")

    # Key findings
    lines.append("\n" + "-" * 50)
    lines.append("KEY FINDINGS")
    lines.append("-" * 50)

    # Find drug with most pancreatitis reports
    panc_counts = df[df['category'] == 'Pancreatitis'].groupby('drug').size()
    if len(panc_counts) > 0:
        max_panc_drug = panc_counts.idxmax()
        max_panc_count = panc_counts.max()
        lines.append(f"\n1. Pancreatitis reports: {max_panc_drug} has the highest count "
                    f"({max_panc_count} reports)")

    # Find drug with most thyroid cancer reports
    thyroid_counts = df[df['category'] == 'Thyroid Cancer'].groupby('drug').size()
    if len(thyroid_counts) > 0:
        max_thy_drug = thyroid_counts.idxmax()
        max_thy_count = thyroid_counts.max()
        lines.append(f"\n2. Thyroid cancer reports: {max_thy_drug} has the highest count "
                    f"({max_thy_count} reports)")

    # Hospitalization rates
    hosp_rates = df.groupby('drug').apply(
        lambda x: (x['is_hospitalization'].sum() / len(x) * 100) if len(x) > 0 else 0
    )
    if len(hosp_rates) > 0:
        max_hosp_drug = hosp_rates.idxmax()
        max_hosp_rate = hosp_rates.max()
        lines.append(f"\n3. Hospitalization rate: {max_hosp_drug} has the highest rate "
                    f"({max_hosp_rate:.1f}%)")

    # Death rates
    death_rates = df.groupby('drug').apply(
        lambda x: (x['is_death'].sum() / len(x) * 100) if len(x) > 0 else 0
    )
    if len(death_rates) > 0 and death_rates.max() > 0:
        max_death_drug = death_rates.idxmax()
        max_death_rate = death_rates.max()
        lines.append(f"\n4. Mortality rate: {max_death_drug} has the highest rate "
                    f"({max_death_rate:.1f}%)")

    lines.append("\n" + "=" * 70)
    lines.append("END OF SUMMARY")
    lines.append("=" * 70)

    summary_text = '\n'.join(lines)

    with open(output_path, 'w') as f:
        f.write(summary_text)

    print(f"  - Saved statistical summary")
    return summary_text


def save_results(df: pd.DataFrame, trends: pd.DataFrame, severity_stats: pd.DataFrame,
                 yoy_changes: dict, output_path: str):
    """Save processed data and analysis results to JSON."""
    print(f"\nSaving analysis results: {output_path}")

    # Calculate summary by drug with proper structure
    drug_summary = {}
    for drug in df['drug'].unique():
        drug_df = df[df['drug'] == drug]
        drug_summary[drug] = {
            'total_reports': len(drug_df),
            'serious_count': int(drug_df['is_serious'].sum()),
            'death_count': int(drug_df['is_death'].sum()),
            'hospitalization_count': int(drug_df['is_hospitalization'].sum())
        }

    # Calculate cross-tabulation with proper structure
    cross_tab = {}
    for drug in df['drug'].unique():
        cross_tab[drug] = {}
        for cat in df['category'].unique():
            count = len(df[(df['drug'] == drug) & (df['category'] == cat)])
            cross_tab[drug][cat] = int(count)

    # Calculate severity by drug with proper structure
    severity_dict = {}
    for drug in df['drug'].unique():
        drug_df = df[df['drug'] == drug]
        severity_dict[drug] = drug_df['severity_category'].value_counts().to_dict()
        # Convert values to int
        severity_dict[drug] = {k: int(v) for k, v in severity_dict[drug].items()}

    # Convert DataFrames to serializable format
    results = {
        'metadata': {
            'analysis_date': datetime.now().isoformat(),
            'total_records_processed': len(df),
            'unique_drugs': int(df['drug'].nunique()),
            'unique_categories': int(df['category'].nunique()),
            'drugs_analyzed': df['drug'].unique().tolist(),
            'categories_analyzed': df['category'].unique().tolist()
        },
        'summary_by_drug': drug_summary,
        'summary_by_category': {k: int(v) for k, v in df.groupby('category').size().to_dict().items()},
        'cross_tabulation': cross_tab,
        'temporal_trends': trends.to_dict(orient='records'),
        'year_over_year_changes': yoy_changes,
        'severity_by_drug': severity_dict,
        'processed_data_sample': df.head(50).to_dict(orient='records')
    }

    with open(output_path, 'w') as f:
        json.dump(results, f, indent=2, default=str)

    print(f"  - Saved analysis results")
    return results


def main():
    """Main execution function."""
    print("=" * 70)
    print("Step 2: Data Preprocessing and Quantitative Signal Analysis")
    print("=" * 70)

    # Step 1: Load data
    faers_path = DATA_DIR / 'faers_raw.json'
    data = load_faers_data(faers_path)
    summary_counts = data.get('summary_counts', {})

    # Step 2: Process records
    df = process_records(data)

    if df.empty:
        print("\nERROR: No valid records processed. Check data quality.")
        return

    # Step 3: Calculate temporal trends
    trends = calculate_temporal_trends(df)

    # Step 4: Calculate severity metrics
    severity_stats = calculate_severity_metrics(df, summary_counts)

    # Step 5: Calculate event counts
    counts_pivot = calculate_event_counts(df)

    # Step 6: Calculate year-over-year changes
    yoy_changes = calculate_year_over_year_change(df)

    # Step 7: Generate visualizations
    plot_time_series(trends, FIGURES_DIR / 'trend_analysis.png')
    plot_severity_analysis(df, FIGURES_DIR / 'severity_analysis.png')
    plot_event_comparison(counts_pivot, FIGURES_DIR / 'event_comparison.png')

    # Step 8: Save results
    results = save_results(df, trends, severity_stats, yoy_changes,
                          RESULTS_DIR / '02_analysis_results.json')

    # Step 9: Generate summary
    summary = generate_summary(df, summary_counts, yoy_changes, severity_stats,
                              RESULTS_DIR / '02_analysis_summary.txt')

    # Print summary to console
    print("\n" + summary)

    print("\n" + "=" * 70)
    print("Analysis Complete!")
    print("=" * 70)
    print(f"\nOutput files:")
    print(f"  - {FIGURES_DIR / 'trend_analysis.png'}")
    print(f"  - {FIGURES_DIR / 'severity_analysis.png'}")
    print(f"  - {FIGURES_DIR / 'event_comparison.png'}")
    print(f"  - {RESULTS_DIR / '02_analysis_results.json'}")
    print(f"  - {RESULTS_DIR / '02_analysis_summary.txt'}")


if __name__ == '__main__':
    main()