#!/usr/bin/env python3
"""
Step 3: Structure-Property Integration
Purpose: Identify polymers in adverse event narratives and merge with chemical properties.

This script:
1. Loads classified events and polymer properties
2. Identifies specific polymers in event text using synonym matching
3. Merges datasets to create Structure-Performance Matrix
4. Generates summary statistics
"""

import pandas as pd
import re
from pathlib import Path
import sys

# Set paths
BASE_DIR = Path("/app/sandbox/session_20251212_210923_d71aa9ce43f6")
CLASSIFIED_EVENTS = BASE_DIR / "data" / "processed" / "classified_events.csv"
POLYMER_PROPERTIES = BASE_DIR / "data" / "raw" / "polymer_properties.csv"
OUTPUT_MATRIX = BASE_DIR / "data" / "processed" / "structure_performance_matrix.csv"
SUMMARY_REPORT = BASE_DIR / "results" / "integration_summary.txt"

# Create directories if needed
OUTPUT_MATRIX.parent.mkdir(parents=True, exist_ok=True)
SUMMARY_REPORT.parent.mkdir(parents=True, exist_ok=True)

print("=" * 80)
print("Step 3: Structure-Property Integration")
print("=" * 80)

# Define polymer synonym mappings (case-insensitive)
POLYMER_SYNONYMS = {
    'PLA': ['pla', 'polylactic', 'polylactide'],
    'PGA': ['pga', 'polyglycolic', 'polyglycolide'],
    'PLGA': ['plga', 'poly(lactic-co-glycolic)', 'polylactide-co-glycolide'],
    'PCL': ['pcl', 'polycaprolactone']
}

def identify_polymer(text):
    """
    Identify polymer type from text using synonym matching.

    Args:
        text: Event narrative text

    Returns:
        str: Polymer identifier ('PLA', 'PGA', 'PLGA', 'PCL', 'Multiple', or 'Unspecified')
    """
    if pd.isna(text):
        return 'Unspecified'

    # Convert to lowercase for case-insensitive matching
    text_lower = text.lower()

    # Track which polymers are found
    found_polymers = set()

    # Search for each polymer's synonyms
    for polymer, synonyms in POLYMER_SYNONYMS.items():
        for synonym in synonyms:
            # Use word boundaries to avoid false positives
            # e.g., "pla" should not match "place" or "plan"
            pattern = r'\b' + re.escape(synonym) + r'\b'
            if re.search(pattern, text_lower):
                found_polymers.add(polymer)
                break  # Found this polymer, move to next

    # Determine result based on findings
    if len(found_polymers) == 0:
        return 'Unspecified'
    elif len(found_polymers) == 1:
        return list(found_polymers)[0]
    else:
        return 'Multiple'

# Step 1: Load Data
print("\n[1/5] Loading data...")
try:
    events_df = pd.read_csv(CLASSIFIED_EVENTS)
    print(f"  ✓ Loaded {len(events_df):,} adverse event records")
    print(f"    Columns: {list(events_df.columns)}")

    properties_df = pd.read_csv(POLYMER_PROPERTIES)
    print(f"  ✓ Loaded {len(properties_df)} polymer property records")
    print(f"    Polymers: {list(properties_df['Polymer'].values)}")
    print(f"    Properties: {list(properties_df.columns)}")
except FileNotFoundError as e:
    print(f"  ✗ Error: Required file not found: {e}")
    sys.exit(1)
except Exception as e:
    print(f"  ✗ Error loading data: {e}")
    sys.exit(1)

# Step 2: Polymer Identification
print("\n[2/5] Identifying polymers in event narratives...")
print("  Using synonym matching with word boundary detection")
print("  Polymer synonyms:")
for polymer, synonyms in POLYMER_SYNONYMS.items():
    print(f"    {polymer}: {', '.join(synonyms)}")

# Apply polymer identification function
events_df['identified_polymer'] = events_df['original_text'].apply(identify_polymer)

# Count polymer identifications
polymer_counts = events_df['identified_polymer'].value_counts()
print(f"\n  Polymer identification results:")
for polymer, count in polymer_counts.items():
    pct = 100 * count / len(events_df)
    print(f"    {polymer:15s}: {count:5,} ({pct:5.1f}%)")

# Step 3: Data Merging
print("\n[3/5] Merging event data with polymer properties...")
# Perform left join to keep all events
# Events with 'Unspecified' or 'Multiple' polymers will have NaN for properties
merged_df = events_df.merge(
    properties_df,
    left_on='identified_polymer',
    right_on='Polymer',
    how='left'
)

print(f"  ✓ Merged dataset shape: {merged_df.shape}")
print(f"    Rows: {len(merged_df):,}")
print(f"    Columns: {len(merged_df.columns)}")

# Check how many records got property data
records_with_properties = merged_df['Polymer'].notna().sum()
pct_with_properties = 100 * records_with_properties / len(merged_df)
print(f"  ✓ Records successfully linked to properties: {records_with_properties:,} ({pct_with_properties:.1f}%)")

# Step 4: Save Output
print("\n[4/5] Saving Structure-Performance Matrix...")
try:
    merged_df.to_csv(OUTPUT_MATRIX, index=False)
    file_size = OUTPUT_MATRIX.stat().st_size / 1024  # KB
    print(f"  ✓ Saved to: {OUTPUT_MATRIX}")
    print(f"    File size: {file_size:.1f} KB")
except Exception as e:
    print(f"  ✗ Error saving output: {e}")
    sys.exit(1)

# Step 5: Generate Summary Report
print("\n[5/5] Generating integration summary report...")

summary_lines = []
summary_lines.append("=" * 80)
summary_lines.append("STRUCTURE-PROPERTY INTEGRATION SUMMARY")
summary_lines.append("=" * 80)
summary_lines.append("")
summary_lines.append("Dataset Overview:")
summary_lines.append(f"  Total adverse event records: {len(events_df):,}")
summary_lines.append(f"  Total polymer types in database: {len(properties_df)}")
summary_lines.append("")
summary_lines.append("Polymer Identification Results:")
summary_lines.append("-" * 80)
for polymer, count in polymer_counts.items():
    pct = 100 * count / len(events_df)
    summary_lines.append(f"  {polymer:15s}: {count:5,} events ({pct:5.1f}%)")
summary_lines.append("")

# Count records linked to specific polymers (not Unspecified or Multiple)
specific_polymers = ['PLA', 'PGA', 'PLGA', 'PCL']
specific_count = events_df['identified_polymer'].isin(specific_polymers).sum()
specific_pct = 100 * specific_count / len(events_df)
summary_lines.append(f"Records linked to specific polymers: {specific_count:,} ({specific_pct:.1f}%)")
summary_lines.append("")

# Failure mode breakdown by polymer
summary_lines.append("Failure Mode Distribution by Polymer:")
summary_lines.append("-" * 80)
failure_cols = ['premature_degradation', 'inflammatory_reaction', 'mechanical_failure', 'incomplete_resorption']

for polymer in specific_polymers:
    if polymer in events_df['identified_polymer'].values:
        polymer_events = events_df[events_df['identified_polymer'] == polymer]
        summary_lines.append(f"\n{polymer} (n={len(polymer_events)}):")
        for col in failure_cols:
            count = polymer_events[col].sum()
            pct = 100 * count / len(polymer_events) if len(polymer_events) > 0 else 0
            summary_lines.append(f"  {col:25s}: {count:4} ({pct:5.1f}%)")

summary_lines.append("")
summary_lines.append("=" * 80)
summary_lines.append("Chemical Properties Added to Dataset:")
summary_lines.append("-" * 80)
summary_lines.append("For each identified polymer, the following properties were merged:")
summary_lines.append("  - Molecular_Weight: Molecular weight (g/mol)")
summary_lines.append("  - XLogP: Octanol-water partition coefficient (lipophilicity)")
summary_lines.append("  - TPSA: Topological polar surface area (Ų)")
summary_lines.append("")

# Display property values for each polymer
summary_lines.append("Property Values by Polymer:")
summary_lines.append("-" * 80)
for _, row in properties_df.iterrows():
    summary_lines.append(f"\n{row['Polymer']} ({row['Full_Name']}):")
    summary_lines.append(f"  Molecular Weight: {row['Molecular_Weight']} g/mol")
    summary_lines.append(f"  XLogP: {row['XLogP']}")
    summary_lines.append(f"  TPSA: {row['TPSA']} Ų")

summary_lines.append("")
summary_lines.append("=" * 80)
summary_lines.append("Output Files:")
summary_lines.append(f"  Structure-Performance Matrix: {OUTPUT_MATRIX}")
summary_lines.append(f"  Integration Summary: {SUMMARY_REPORT}")
summary_lines.append("=" * 80)

# Write summary to file
summary_text = "\n".join(summary_lines)
SUMMARY_REPORT.write_text(summary_text)
print(f"  ✓ Saved to: {SUMMARY_REPORT}")

# Also print summary to console
print("\n" + summary_text)

print("\n" + "=" * 80)
print("✓ Step 3 Complete: Structure-Property Integration")
print("=" * 80)
print(f"\nOutputs:")
print(f"  1. {OUTPUT_MATRIX}")
print(f"  2. {SUMMARY_REPORT}")
print("\nKey Findings:")
print(f"  • {specific_count:,} events ({specific_pct:.1f}%) were successfully linked to specific polymers")
print(f"  • {records_with_properties:,} events now have chemical property data")
print(f"  • Ready for Step 4: Structure-Property Correlation Analysis")