"""
Example: Fetching Heracles Reports

This example demonstrates how to fetch various Heracles characterization
reports for a cohort.

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

from smart_omop.client import OMOPClient
import json
from pathlib import Path


def main():
    # Configuration
    BASE_URL = "http://your-server:8080/WebAPI"
    COHORT_ID = 164
    SOURCE_KEY = "KAGGLECOPD"

    # Create client
    with OMOPClient(BASE_URL) as client:

        print("Fetching Heracles reports for cohort", COHORT_ID)
        print("-" * 60)

        # 1. Fetch person/demographic report
        print("\n1. Person Demographics Report")
        try:
            person_report = client.get_heracles_person_report(
                COHORT_ID,
                SOURCE_KEY,
                refresh=True
            )
            print(f"   ✓ Retrieved {len(person_report)} analyses")

            # Example: Extract person count
            if person_report:
                person_count_analysis = next(
                    (a for a in person_report if a.get('analysisId') == 1),
                    None
                )
                if person_count_analysis:
                    count = person_count_analysis.get('countValue', 0)
                    print(f"   Total persons: {count}")

        except Exception as e:
            print(f"   ✗ Error: {e}")

        # 2. Fetch condition occurrence report
        print("\n2. Condition Occurrence Report")
        try:
            condition_report = client.get_heracles_condition_report(
                COHORT_ID,
                SOURCE_KEY,
                refresh=True
            )
            print(f"   ✓ Retrieved {len(condition_report)} analyses")

            # Example: Top 10 conditions
            if condition_report:
                # Analysis 400: Most common conditions
                top_conditions = [
                    a for a in condition_report
                    if a.get('analysisId') == 400
                ][:10]
                if top_conditions:
                    print(f"   Top {len(top_conditions)} conditions:")
                    for cond in top_conditions:
                        concept_id = cond.get('stratum1')
                        count = cond.get('countValue', 0)
                        print(f"     - Concept {concept_id}: {count} occurrences")

        except Exception as e:
            print(f"   ✗ Error: {e}")

        # 3. Fetch drug exposure report
        print("\n3. Drug Exposure Report")
        try:
            drug_report = client.get_heracles_drug_report(
                COHORT_ID,
                SOURCE_KEY,
                refresh=True
            )
            print(f"   ✓ Retrieved {len(drug_report)} analyses")

        except Exception as e:
            print(f"   ✗ Error: {e}")

        # 4. Fetch procedure occurrence report
        print("\n4. Procedure Occurrence Report")
        try:
            procedure_report = client.get_heracles_procedure_report(
                COHORT_ID,
                SOURCE_KEY,
                refresh=True
            )
            print(f"   ✓ Retrieved {len(procedure_report)} analyses")

        except Exception as e:
            print(f"   ✗ Error: {e}")

        # 5. Fetch measurement report
        print("\n5. Measurement Report")
        try:
            measurement_report = client.get_heracles_measurement_report(
                COHORT_ID,
                SOURCE_KEY,
                refresh=True
            )
            print(f"   ✓ Retrieved {len(measurement_report)} analyses")

        except Exception as e:
            print(f"   ✗ Error: {e}")

        # 6. Fetch dashboard summary
        print("\n6. Dashboard Summary")
        try:
            dashboard = client.get_heracles_dashboard_report(
                COHORT_ID,
                SOURCE_KEY,
                refresh=True
            )
            print(f"   ✓ Retrieved dashboard with {len(dashboard)} sections")

        except Exception as e:
            print(f"   ✗ Error: {e}")

        # 7. Export all reports to files
        print("\n7. Exporting All Reports to Files")
        print("-" * 60)

        output_dir = Path(f"heracles_reports_{COHORT_ID}")
        output_dir.mkdir(exist_ok=True)

        report_types = {
            'person': client.get_heracles_person_report,
            'condition': client.get_heracles_condition_report,
            'drug': client.get_heracles_drug_report,
            'procedure': client.get_heracles_procedure_report,
            'measurement': client.get_heracles_measurement_report,
            'observation': client.get_heracles_observation_report,
            'death': client.get_heracles_death_report,
            'dashboard': client.get_heracles_dashboard_report,
        }

        for report_name, fetch_func in report_types.items():
            try:
                print(f"   Exporting {report_name}...", end=' ')
                report_data = fetch_func(COHORT_ID, SOURCE_KEY, refresh=True)

                filename = output_dir / f"{report_name}_report.json"
                with open(filename, 'w') as f:
                    json.dump(report_data, f, indent=2)

                print(f"✓ {filename}")

            except Exception as e:
                print(f"✗ Error: {e}")

        print(f"\n✓ All reports exported to: {output_dir}")

        # 8. Using the generic method for custom report types
        print("\n8. Generic Method for Custom Reports")
        print("-" * 60)
        try:
            # Fetch any report type using the generic method
            custom_report = client.get_heracles_report(
                COHORT_ID,
                SOURCE_KEY,
                "components_summary",  # Custom report type
                refresh=True
            )
            print(f"   ✓ Retrieved custom report with {len(custom_report)} items")

        except Exception as e:
            print(f"   ✗ Error: {e}")


if __name__ == "__main__":
    main()
