#!/usr/bin/env python3
"""
Step 2: Failure Mode Classification
Process FDA adverse event narratives to categorize failure modes and polymer-relevance.
"""

import json
import pandas as pd
import re
from pathlib import Path
from datetime import datetime

# Set paths
BASE_DIR = Path("/app/sandbox/session_20251212_210923_d71aa9ce43f6")
RAW_DATA = BASE_DIR / "data/raw/fda_maude_events.json"
PROCESSED_DIR = BASE_DIR / "data/processed"
RESULTS_DIR = BASE_DIR / "results"

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

# Define keyword categories (case-insensitive matching)
KEYWORD_CATEGORIES = {
    "polymer_relevant": [
        "polymer", "coating", "absorbable", "resorbable", "biodegradable",
        "scaffold", "polylactide", "plga", "pla"
    ],
    "premature_degradation": [
        "premature", "degraded early", "fragmented", "disintegrated", "particle"
    ],
    "inflammatory_reaction": [
        "inflammation", "swelling", "edema", "erythema", "reaction",
        "immune", "hypersensitivity"
    ],
    "mechanical_failure": [
        "fracture", "break", "snap", "crack", "structural failure", "collapse"
    ],
    "incomplete_resorption": [
        "remained", "persistent", "not absorbed", "encapsulated", "residual"
    ]
}


def extract_narrative_text(mdr_text_list):
    """
    Extract and combine all narrative text from the mdr_text list.

    Args:
        mdr_text_list: List of dictionaries containing text fields

    Returns:
        Combined narrative text as a single string
    """
    if not mdr_text_list or not isinstance(mdr_text_list, list):
        return ""

    text_parts = []
    for entry in mdr_text_list:
        if isinstance(entry, dict) and 'text' in entry:
            text_parts.append(entry['text'])

    return " ".join(text_parts)


def preprocess_text(text):
    """
    Clean and preprocess text for keyword matching.

    Args:
        text: Raw text string

    Returns:
        Cleaned text in lowercase
    """
    if not text:
        return ""

    # Convert to lowercase
    text = text.lower()

    # Remove special characters but keep spaces and basic punctuation
    # This preserves multi-word phrases like "degraded early"
    text = re.sub(r'[^a-z0-9\s\-]', ' ', text)

    # Normalize whitespace
    text = re.sub(r'\s+', ' ', text).strip()

    return text


def classify_text(text, keywords):
    """
    Check if any keywords appear in the text.

    Args:
        text: Preprocessed text string
        keywords: List of keywords to search for

    Returns:
        1 if any keyword found, 0 otherwise
    """
    if not text:
        return 0

    for keyword in keywords:
        # Use word boundary matching for single words
        # For phrases, check for the exact phrase
        if ' ' in keyword:
            # Multi-word phrase - check for exact phrase
            if keyword in text:
                return 1
        else:
            # Single word - use word boundary
            pattern = r'\b' + re.escape(keyword) + r'\b'
            if re.search(pattern, text):
                return 1

    return 0


def main():
    print("=" * 60)
    print("Step 2: Failure Mode Classification")
    print("=" * 60)

    # Load raw data
    print(f"\n[1/6] Loading data from {RAW_DATA}...")
    with open(RAW_DATA, 'r') as f:
        raw_data = json.load(f)
    print(f"✓ Loaded {len(raw_data)} records")

    # Extract and process text
    print("\n[2/6] Extracting narrative text from records...")
    processed_records = []

    for i, record in enumerate(raw_data):
        if i % 50 == 0:
            print(f"  Processing record {i}/{len(raw_data)}...")

        # Extract report metadata
        report_number = record.get('report_number', f'UNKNOWN_{i}')
        date_received = record.get('date_received', '')

        # Extract and combine narrative text
        mdr_text_list = record.get('mdr_text', [])
        raw_text = extract_narrative_text(mdr_text_list)

        # Preprocess text
        cleaned_text = preprocess_text(raw_text)

        # Truncate for storage (keep first 1000 chars of original)
        truncated_text = raw_text[:1000] if raw_text else ""

        processed_records.append({
            'report_number': report_number,
            'date_received': date_received,
            'original_text': truncated_text,
            'preprocessed_text': cleaned_text
        })

    print(f"✓ Extracted text from {len(processed_records)} records")

    # Apply classification
    print("\n[3/6] Applying keyword-based classification...")
    df = pd.DataFrame(processed_records)

    for category, keywords in KEYWORD_CATEGORIES.items():
        print(f"  Classifying: {category} ({len(keywords)} keywords)")
        df[category] = df['preprocessed_text'].apply(
            lambda text: classify_text(text, keywords)
        )

    print("✓ Classification complete")

    # Generate summary statistics
    print("\n[4/6] Generating classification summary...")
    summary_stats = {}
    for category in KEYWORD_CATEGORIES.keys():
        count = df[category].sum()
        percentage = (count / len(df) * 100)
        summary_stats[category] = {
            'count': int(count),
            'percentage': round(percentage, 2)
        }
        print(f"  {category}: {count} reports ({percentage:.1f}%)")

    # Check for reports with multiple failure modes
    failure_mode_cols = [
        'premature_degradation', 'inflammatory_reaction',
        'mechanical_failure', 'incomplete_resorption'
    ]
    df['failure_mode_count'] = df[failure_mode_cols].sum(axis=1)
    multi_failure = (df['failure_mode_count'] > 1).sum()
    print(f"\n  Reports with multiple failure modes: {multi_failure}")

    # Polymer-relevant subset
    polymer_relevant = df['polymer_relevant'].sum()
    print(f"  Polymer-relevant reports: {polymer_relevant} ({polymer_relevant/len(df)*100:.1f}%)")

    # Save classified data
    print("\n[5/6] Saving classified data...")
    output_file = PROCESSED_DIR / "classified_events.csv"

    # Select columns for output (exclude preprocessed_text to save space)
    output_cols = ['report_number', 'date_received', 'original_text'] + list(KEYWORD_CATEGORIES.keys())
    df_output = df[output_cols]

    df_output.to_csv(output_file, index=False)
    print(f"✓ Saved to {output_file}")
    print(f"  Shape: {df_output.shape}")
    print(f"  Size: {output_file.stat().st_size / 1024:.1f} KB")

    # Save summary report
    print("\n[6/6] Saving summary report...")
    summary_file = RESULTS_DIR / "failure_mode_counts.txt"

    with open(summary_file, 'w') as f:
        f.write("=" * 70 + "\n")
        f.write("FDA MAUDE Adverse Event Classification Summary\n")
        f.write("Step 2: Failure Mode Classification\n")
        f.write("=" * 70 + "\n")
        f.write(f"\nGenerated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
        f.write(f"Total Reports Analyzed: {len(df)}\n")
        f.write("\n" + "-" * 70 + "\n")
        f.write("CLASSIFICATION RESULTS\n")
        f.write("-" * 70 + "\n\n")

        for category, stats in summary_stats.items():
            category_name = category.replace('_', ' ').title()
            f.write(f"{category_name}:\n")
            f.write(f"  Count: {stats['count']}\n")
            f.write(f"  Percentage: {stats['percentage']}%\n\n")

        f.write("-" * 70 + "\n")
        f.write("ADDITIONAL INSIGHTS\n")
        f.write("-" * 70 + "\n\n")
        f.write(f"Reports with Multiple Failure Modes: {multi_failure}\n")
        f.write(f"Polymer-Relevant Reports: {polymer_relevant} ({polymer_relevant/len(df)*100:.1f}%)\n")
        f.write(f"\nPolymer-Relevant AND Failure Mode Breakdown:\n")

        for mode in failure_mode_cols:
            polymer_and_failure = df[(df['polymer_relevant'] == 1) & (df[mode] == 1)].shape[0]
            mode_name = mode.replace('_', ' ').title()
            f.write(f"  {mode_name}: {polymer_and_failure}\n")

        f.write("\n" + "=" * 70 + "\n")

    print(f"✓ Saved summary to {summary_file}")

    # Display summary to console
    print("\n" + "=" * 60)
    print("CLASSIFICATION SUMMARY")
    print("=" * 60)
    with open(summary_file, 'r') as f:
        print(f.read())

    print("\n" + "=" * 60)
    print("✓ Step 2 Complete!")
    print("=" * 60)
    print(f"\nOutputs:")
    print(f"  - {output_file}")
    print(f"  - {summary_file}")
    print("\nNext: Step 3 - Quantitative Analysis")


if __name__ == "__main__":
    main()