"""
Visualizations for OMOP Cohort Characterization

Visualizations for cohort analysis:
- Demographics distributions (age pyramids, gender breakdowns)
- Condition prevalence treemaps
- Temporal patterns (cohort entry over time)
- Measurement distributions

Uses Plotly for interactive charts with matplotlib fallback.

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

from __future__ import annotations
from pathlib import Path
from typing import Dict, Any, List, Optional, Tuple
import json


class CohortVisualizer:
    """
    Visualizations for OMOP cohort characterization.

    Creates interactive plots for cohort analysis.
    """

    def __init__(self, output_dir: str = "./visualizations"):
        """
        Initialize visualizer.

        Args:
            output_dir: Directory to save visualization outputs
        """
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)

        # Color schemes
        self.colors = {
            'primary': '#005EB8',     # NHS Medical Blue
            'secondary': '#003087',    # Dark Blue
            'accent': '#41B6E6',       # Light Blue
            'success': '#007F3B',      # Green
            'warning': '#FFB81C',      # Amber
            'danger': '#DA291C',       # Red
            'neutral': '#768692'       # Grey
        }

    def create_age_pyramid(
        self,
        age_data: Dict[str, List[int]],
        save_path: Optional[str] = None
    ) -> str:
        """
        Create age pyramid visualization by gender.

        Args:
            age_data: Dict with 'male' and 'female' keys containing age lists
            save_path: Optional custom save path

        Returns:
            Path to saved visualization
        """
        try:
            import plotly.graph_objects as go
            from plotly.subplots import make_subplots
        except ImportError:
            return self._create_fallback_age_chart(age_data, save_path)

        # Bin ages into 5-year groups
        bins = list(range(0, 105, 5))
        bin_labels = [f"{i}-{i+4}" for i in bins[:-1]]

        male_counts = self._bin_ages(age_data.get('male', []), bins)
        female_counts = self._bin_ages(age_data.get('female', []), bins)

        # Create pyramid (males negative, females positive)
        fig = go.Figure()

        fig.add_trace(go.Bar(
            y=bin_labels,
            x=[-c for c in male_counts],
            name='Male',
            orientation='h',
            marker=dict(color=self.colors['primary'])
        ))

        fig.add_trace(go.Bar(
            y=bin_labels,
            x=female_counts,
            name='Female',
            orientation='h',
            marker=dict(color=self.colors['accent'])
        ))

        fig.update_layout(
            title='Cohort Age Distribution by Gender',
            barmode='overlay',
            bargap=0.1,
            xaxis=dict(title='Population Count'),
            yaxis=dict(title='Age Group'),
            template='plotly_white',
            height=600
        )

        # Save
        if not save_path:
            save_path = str(self.output_dir / 'age_pyramid.html')

        fig.write_html(save_path)
        return save_path

    def create_condition_treemap(
        self,
        condition_counts: Dict[str, int],
        condition_names: Optional[Dict[str, str]] = None,
        save_path: Optional[str] = None
    ) -> str:
        """
        Create treemap visualization of condition prevalence.

        Args:
            condition_counts: Dict of concept_id -> count
            condition_names: Optional dict of concept_id -> name
            save_path: Optional custom save path

        Returns:
            Path to saved visualization
        """
        try:
            import plotly.express as px
        except ImportError:
            return self._create_fallback_bar_chart(condition_counts, save_path)

        # Prepare data
        names = condition_names or {}
        labels = []
        values = []
        ids = []

        for concept_id, count in condition_counts.items():
            ids.append(concept_id)
            labels.append(names.get(concept_id, f"Concept {concept_id}"))
            values.append(count)

        # Create treemap
        fig = px.treemap(
            names=labels,
            values=values,
            title='Condition Prevalence in Cohort',
            color_discrete_sequence=[self.colors['primary'], self.colors['secondary'],
                                    self.colors['accent'], self.colors['success']]
        )

        fig.update_layout(
            template='plotly_white',
            height=600
        )

        # Save
        if not save_path:
            save_path = str(self.output_dir / 'condition_treemap.html')

        fig.write_html(save_path)
        return save_path

    def create_temporal_pattern(
        self,
        dates: List[str],
        save_path: Optional[str] = None
    ) -> str:
        """
        Create temporal pattern visualization (cohort entry over time).

        Args:
            dates: List of dates (YYYY-MM-DD format)
            save_path: Optional custom save path

        Returns:
            Path to saved visualization
        """
        try:
            import plotly.graph_objects as go
            from collections import Counter
            from datetime import datetime
        except ImportError:
            return self._create_fallback_temporal(dates, save_path)

        # Group by month
        month_counts = Counter()
        for date_str in dates:
            try:
                dt = datetime.strptime(date_str, '%Y-%m-%d')
                month_key = dt.strftime('%Y-%m')
                month_counts[month_key] += 1
            except:
                continue

        # Sort by date
        sorted_months = sorted(month_counts.items())
        months = [m[0] for m in sorted_months]
        counts = [m[1] for m in sorted_months]

        # Create line chart
        fig = go.Figure()

        fig.add_trace(go.Scatter(
            x=months,
            y=counts,
            mode='lines+markers',
            name='Cohort Entry',
            line=dict(color=self.colors['primary'], width=2),
            marker=dict(size=6)
        ))

        fig.update_layout(
            title='Cohort Entry Over Time',
            xaxis=dict(title='Month'),
            yaxis=dict(title='New Cohort Members'),
            template='plotly_white',
            height=400
        )

        # Save
        if not save_path:
            save_path = str(self.output_dir / 'temporal_pattern.html')

        fig.write_html(save_path)
        return save_path

    def create_dashboard(
        self,
        cohort_data: Dict[str, Any],
        save_path: Optional[str] = None
    ) -> str:
        """
        Create comprehensive cohort characterization dashboard.

        Args:
            cohort_data: Cohort data dict from fetch_cohort_data
            save_path: Optional custom save path

        Returns:
            Path to saved dashboard
        """
        try:
            import plotly.graph_objects as go
            from plotly.subplots import make_subplots
        except ImportError:
            return self._create_fallback_dashboard(cohort_data, save_path)

        # Extract key metrics
        person_count = cohort_data.get('results', {}).get('personCount', 0)
        status = cohort_data.get('results', {}).get('status', 'N/A')

        # Create dashboard with subplots
        fig = make_subplots(
            rows=2, cols=2,
            subplot_titles=('Cohort Summary', 'Demographics', 'Status', 'Metrics'),
            specs=[[{'type': 'indicator'}, {'type': 'bar'}],
                   [{'type': 'indicator'}, {'type': 'table'}]]
        )

        # Person count indicator
        fig.add_trace(go.Indicator(
            mode='number',
            value=person_count,
            title={'text': 'Total Persons'},
            number={'font': {'size': 48, 'color': self.colors['primary']}}
        ), row=1, col=1)

        # Status indicator
        status_color = self.colors['success'] if status == 'COMPLETE' else self.colors['warning']
        fig.add_trace(go.Indicator(
            mode='text',
            value=0,
            title={'text': 'Generation Status'},
            number={'font': {'size': 24, 'color': status_color}},
            domain={'x': [0, 1], 'y': [0, 1]}
        ), row=2, col=1)

        fig.update_layout(
            title='Cohort Characterization Dashboard',
            template='plotly_white',
            height=800
        )

        # Save
        if not save_path:
            save_path = str(self.output_dir / 'cohort_dashboard.html')

        fig.write_html(save_path)
        return save_path

    def _bin_ages(self, ages: List[int], bins: List[int]) -> List[int]:
        """Bin ages into histogram groups."""
        counts = [0] * (len(bins) - 1)

        for age in ages:
            for i in range(len(bins) - 1):
                if bins[i] <= age < bins[i + 1]:
                    counts[i] += 1
                    break

        return counts

    def _create_fallback_age_chart(
        self,
        age_data: Dict[str, List[int]],
        save_path: Optional[str]
    ) -> str:
        """Fallback to matplotlib if plotly not available."""
        try:
            import matplotlib.pyplot as plt
        except ImportError:
            return str(self.output_dir / 'visualization_not_available.txt')

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

        male_ages = age_data.get('male', [])
        female_ages = age_data.get('female', [])

        ax.hist([male_ages, female_ages], bins=20, label=['Male', 'Female'],
                color=[self.colors['primary'], self.colors['accent']])

        ax.set_xlabel('Age')
        ax.set_ylabel('Count')
        ax.set_title('Age Distribution by Gender')
        ax.legend()

        if not save_path:
            save_path = str(self.output_dir / 'age_distribution.png')

        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.close()

        return save_path

    def _create_fallback_bar_chart(
        self,
        data: Dict[str, int],
        save_path: Optional[str]
    ) -> str:
        """Fallback bar chart for condition data."""
        try:
            import matplotlib.pyplot as plt
        except ImportError:
            return str(self.output_dir / 'visualization_not_available.txt')

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

        labels = list(data.keys())[:10]  # Top 10
        values = [data[k] for k in labels]

        ax.bar(labels, values, color=self.colors['primary'])
        ax.set_xlabel('Condition Concept ID')
        ax.set_ylabel('Count')
        ax.set_title('Top 10 Conditions')
        ax.tick_params(axis='x', rotation=45)

        if not save_path:
            save_path = str(self.output_dir / 'conditions_bar.png')

        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.close()

        return save_path

    def _create_fallback_temporal(self, dates: List[str], save_path: Optional[str]) -> str:
        """Fallback temporal visualization."""
        save_path = save_path or str(self.output_dir / 'temporal_fallback.json')

        with open(save_path, 'w') as f:
            json.dump({'dates': dates, 'count': len(dates)}, f, indent=2)

        return save_path

    def _create_fallback_dashboard(
        self,
        cohort_data: Dict[str, Any],
        save_path: Optional[str]
    ) -> str:
        """Fallback dashboard as JSON."""
        save_path = save_path or str(self.output_dir / 'dashboard_fallback.json')

        with open(save_path, 'w') as f:
            json.dump(cohort_data, f, indent=2)

        return save_path

    def create_age_distribution(
        self,
        person_report: Dict[str, Any],
        save_path: Optional[str] = None
    ) -> str:
        """
        Create age distribution from Heracles person report.

        Args:
            person_report: Heracles person report data
            save_path: Optional custom save path

        Returns:
            Path to saved visualization
        """
        try:
            import plotly.graph_objects as go
        except ImportError:
            save_path = save_path or str(self.output_dir / 'age_dist_fallback.json')
            with open(save_path, 'w') as f:
                json.dump(person_report.get('yearOfBirth', []), f, indent=2)
            return save_path

        # Extract year of birth data
        birth_data = person_report.get('yearOfBirth', [])
        stats = person_report.get('yearOfBirthStats', [{}])[0]

        intervals = [d['intervalIndex'] for d in birth_data]
        counts = [d['countValue'] for d in birth_data]

        fig = go.Figure(data=[
            go.Bar(x=intervals, y=counts, marker=dict(color=self.colors['primary']))
        ])

        min_year = stats.get('minValue', 'N/A')
        max_year = stats.get('maxValue', 'N/A')

        fig.update_layout(
            title=f'Birth Year Distribution ({min_year}-{max_year})',
            xaxis_title='Year Interval Index',
            yaxis_title='Number of Persons',
            template='plotly_white'
        )

        if not save_path:
            save_path = str(self.output_dir / 'age_distribution.html')

        fig.write_html(save_path)
        return save_path

    def create_gender_distribution(
        self,
        person_report: Dict[str, Any],
        save_path: Optional[str] = None
    ) -> str:
        """
        Create gender distribution from Heracles person report.

        Args:
            person_report: Heracles person report data
            save_path: Optional custom save path

        Returns:
            Path to saved visualization
        """
        try:
            import plotly.graph_objects as go
        except ImportError:
            save_path = save_path or str(self.output_dir / 'gender_fallback.json')
            with open(save_path, 'w') as f:
                json.dump(person_report.get('gender', []), f, indent=2)
            return save_path

        gender_data = person_report.get('gender', [])

        labels = [g['conceptName'] for g in gender_data]
        values = [g['countValue'] for g in gender_data]

        fig = go.Figure(data=[
            go.Pie(
                labels=labels,
                values=values,
                marker=dict(colors=[self.colors['primary'], self.colors['accent']]),
                hole=0.3
            )
        ])

        fig.update_layout(
            title='Gender Distribution',
            template='plotly_white'
        )

        if not save_path:
            save_path = str(self.output_dir / 'gender_distribution.html')

        fig.write_html(save_path)
        return save_path

    def create_condition_prevalence(
        self,
        condition_report: List[Dict[str, Any]],
        save_path: Optional[str] = None,
        top_n: int = 10
    ) -> str:
        """
        Create condition prevalence treemap from Heracles condition report.

        Args:
            condition_report: Heracles condition report data
            save_path: Optional custom save path
            top_n: Number of top conditions to display

        Returns:
            Path to saved visualization
        """
        try:
            import plotly.graph_objects as go
        except ImportError:
            save_path = save_path or str(self.output_dir / 'conditions_fallback.json')
            with open(save_path, 'w') as f:
                json.dump(condition_report[:top_n], f, indent=2)
            return save_path

        # Sort by prevalence and take top N
        sorted_conditions = sorted(
            condition_report,
            key=lambda x: x.get('numPersons', 0),
            reverse=True
        )[:top_n]

        labels = []
        parents = []
        values = []

        # Root
        labels.append("All Conditions")
        parents.append("")
        values.append(0)

        for cond in sorted_conditions:
            name = cond.get('conceptPath', '').split('||')[-1]
            if not name or name == 'NA':
                name = f"Concept {cond.get('conceptId', 'Unknown')}"

            num_persons = cond.get('numPersons', 0)
            percent = cond.get('percentPersons', 0)

            label = f"{name}<br>{num_persons} persons ({percent:.1%})"

            labels.append(label)
            parents.append("All Conditions")
            values.append(num_persons)

        fig = go.Figure(go.Treemap(
            labels=labels,
            parents=parents,
            values=values,
            marker=dict(colorscale='Blues')
        ))

        fig.update_layout(
            title='Condition Prevalence',
            template='plotly_white'
        )

        if not save_path:
            save_path = str(self.output_dir / 'condition_prevalence.html')

        fig.write_html(save_path)
        return save_path

    def create_dashboard_from_reports(
        self,
        reports: Dict[str, Any],
        save_path: Optional[str] = None
    ) -> str:
        """
        Create dashboard from multiple Heracles reports.

        Args:
            reports: Dict with 'person', 'condition', etc. reports
            save_path: Optional custom save path

        Returns:
            Path to saved visualization
        """
        try:
            from plotly.subplots import make_subplots
            import plotly.graph_objects as go
        except ImportError:
            save_path = save_path or str(self.output_dir / 'dashboard_fallback.json')
            with open(save_path, 'w') as f:
                json.dump({'reports': list(reports.keys())}, f, indent=2)
            return save_path

        person_data = reports.get('person', {})
        condition_data = reports.get('condition', [])

        # Create subplots
        fig = make_subplots(
            rows=2, cols=2,
            subplot_titles=('Gender Distribution', 'Birth Year Distribution',
                          'Top Conditions', 'Summary'),
            specs=[[{'type': 'pie'}, {'type': 'bar'}],
                   [{'type': 'bar'}, {'type': 'table'}]]
        )

        # Gender pie chart
        gender_data = person_data.get('gender', [])
        if gender_data:
            fig.add_trace(go.Pie(
                labels=[g['conceptName'] for g in gender_data],
                values=[g['countValue'] for g in gender_data],
                marker=dict(colors=[self.colors['primary'], self.colors['accent']])
            ), row=1, col=1)

        # Birth year bar chart
        birth_data = person_data.get('yearOfBirth', [])
        if birth_data:
            fig.add_trace(go.Bar(
                x=[d['intervalIndex'] for d in birth_data[:20]],
                y=[d['countValue'] for d in birth_data[:20]],
                marker=dict(color=self.colors['primary'])
            ), row=1, col=2)

        # Top conditions
        if condition_data:
            top_5 = sorted(condition_data, key=lambda x: x.get('numPersons', 0), reverse=True)[:5]
            names = [c.get('conceptPath', '').split('||')[-1][:20] for c in top_5]
            counts = [c.get('numPersons', 0) for c in top_5]

            fig.add_trace(go.Bar(
                x=counts,
                y=names,
                orientation='h',
                marker=dict(color=self.colors['accent'])
            ), row=2, col=1)

        # Summary table
        total_persons = sum(g['countValue'] for g in gender_data) if gender_data else 0
        total_conditions = len(condition_data)

        fig.add_trace(go.Table(
            header=dict(values=['Metric', 'Value']),
            cells=dict(values=[
                ['Total Persons', 'Total Conditions', 'Gender Ratio (M/F)'],
                [total_persons, total_conditions,
                 f"{gender_data[0]['countValue']}/{gender_data[1]['countValue']}" if len(gender_data) == 2 else 'N/A']
            ])
        ), row=2, col=2)

        fig.update_layout(
            title_text="Cohort Characterization Dashboard",
            showlegend=False,
            height=800,
            template='plotly_white'
        )

        if not save_path:
            save_path = str(self.output_dir / 'dashboard.html')

        fig.write_html(save_path)
        return save_path


def create_cohort_visualizations(
    cohort_data: Dict[str, Any],
    output_dir: str = "./visualizations"
) -> Dict[str, str]:
    """
    Create all standard visualizations for cohort data.

    Args:
        cohort_data: Cohort data from fetch_cohort_data
        output_dir: Output directory for visualizations

    Returns:
        Dict mapping visualization type to file path
    """
    visualizer = CohortVisualizer(output_dir)

    outputs = {}

    # Dashboard
    outputs['dashboard'] = visualizer.create_dashboard(cohort_data)

    return outputs
