"""
Example: Creating Visualizations from Heracles Reports

Demonstrates how to create visualizations from Heracles characterization reports.

Author: Ankur Lohachab
Department of Advanced Computing Sciences, Maastricht University
"""

import json
from smart_omop import CohortVisualizer, OMOPClient


def visualize_from_files():
    """Create visualizations from exported Heracles report files."""

    # Load reports from files
    with open('cohort167_KAGGLECOPD_person.json') as f:
        person_data = json.load(f)

    with open('cohort167_KAGGLECOPD_condition.json') as f:
        condition_data = json.load(f)

    # Initialize visualizer
    visualizer = CohortVisualizer(output_dir="./visualizations")

    # Create age distribution
    print("Creating age distribution...")
    age_path = visualizer.create_age_distribution(person_data)
    print(f"  Saved to: {age_path}")

    # Create gender distribution
    print("Creating gender distribution...")
    gender_path = visualizer.create_gender_distribution(person_data)
    print(f"  Saved to: {gender_path}")

    # Create condition prevalence treemap
    print("Creating condition prevalence treemap...")
    condition_path = visualizer.create_condition_prevalence(condition_data, top_n=10)
    print(f"  Saved to: {condition_path}")

    # Create dashboard with all charts
    print("Creating dashboard...")
    dashboard_path = visualizer.create_dashboard_from_reports({
        'person': person_data,
        'condition': condition_data
    })
    print(f"  Saved to: {dashboard_path}")


def visualize_from_api():
    """Fetch reports via API and create visualizations."""

    BASE_URL = "http://your-server:8080/WebAPI"
    COHORT_ID = 167
    SOURCE_KEY = "KAGGLECOPD"

    with OMOPClient(BASE_URL) as client:
        print(f"Fetching reports for cohort {COHORT_ID}...")

        # Fetch reports
        person_report = client.get_heracles_person_report(
            COHORT_ID,
            SOURCE_KEY,
            refresh=True
        )

        condition_report = client.get_heracles_condition_report(
            COHORT_ID,
            SOURCE_KEY,
            refresh=True
        )

        # Create visualizations
        visualizer = CohortVisualizer(output_dir="./visualizations")

        # Age distribution
        print("Creating age distribution...")
        age_path = visualizer.create_age_distribution(person_report)
        print(f"  Saved to: {age_path}")

        # Gender pie chart
        print("Creating gender distribution...")
        gender_path = visualizer.create_gender_distribution(person_report)
        print(f"  Saved to: {gender_path}")

        # Condition treemap
        print("Creating condition prevalence...")
        condition_path = visualizer.create_condition_prevalence(condition_report)
        print(f"  Saved to: {condition_path}")

        # Dashboard
        print("Creating dashboard...")
        dashboard_path = visualizer.create_dashboard_from_reports({
            'person': person_report,
            'condition': condition_report
        })
        print(f"  Saved to: {dashboard_path}")


def parse_gender_data(person_report):
    """Parse gender data from person report."""
    gender_data = person_report.get('gender', [])

    print("\nGender Distribution:")
    for gender in gender_data:
        name = gender['conceptName']
        count = gender['countValue']
        print(f"  {name}: {count} persons")


def parse_birth_year_data(person_report):
    """Parse birth year data from person report."""
    birth_data = person_report.get('yearOfBirth', [])
    stats = person_report.get('yearOfBirthStats', [{}])[0]

    print(f"\nBirth Year Statistics:")
    print(f"  Range: {stats.get('minValue')} to {stats.get('maxValue')}")
    print(f"  Total intervals: {len(birth_data)}")

    # Show top 5 intervals
    top_5 = sorted(birth_data, key=lambda x: x['countValue'], reverse=True)[:5]
    print(f"\n  Top 5 birth year intervals:")
    for entry in top_5:
        idx = entry['intervalIndex']
        count = entry['countValue']
        percent = entry['percentValue']
        print(f"    Interval {idx}: {count} persons ({percent:.1%})")


def parse_condition_data(condition_report):
    """Parse condition data from condition report."""
    print(f"\nCondition Analysis:")
    print(f"  Total unique conditions: {len(condition_report)}")

    # Top 5 by prevalence
    top_5 = sorted(condition_report, key=lambda x: x.get('numPersons', 0), reverse=True)[:5]

    print(f"\n  Top 5 conditions by prevalence:")
    for cond in top_5:
        name = cond.get('conceptPath', '').split('||')[-1]
        concept_id = cond.get('conceptId')
        num_persons = cond.get('numPersons', 0)
        percent = cond.get('percentPersons', 0)
        records_per_person = cond.get('recordsPerPerson', 0)

        print(f"\n    {name} ({concept_id})")
        print(f"      Persons: {num_persons} ({percent:.1%})")
        print(f"      Records per person: {records_per_person:.2f}")


if __name__ == "__main__":
    import sys

    if len(sys.argv) > 1 and sys.argv[1] == "--from-api":
        # Fetch from API and visualize
        visualize_from_api()
    else:
        # Use local files
        print("Using local report files...")
        print("(Use --from-api to fetch from WebAPI)")
        print()

        try:
            visualize_from_files()
        except FileNotFoundError as e:
            print(f"\nError: {e}")
            print("\nFirst export reports using:")
            print("smart-omop --base-url <URL> export-reports \\")
            print("  --cohort-id 167 --source-key KAGGLECOPD \\")
            print("  --output-dir . --refresh")
