from django.shortcuts import render # type: ignore

# Create your views here.
import io
import base64
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from django.shortcuts import render # type: ignore
from lifelines import KaplanMeierFitter # type: ignore
from .models import ClinicalPatient, TreatmentTimeline
import plotly.graph_objs as go # type: ignore
from plotly.io import to_html
import pandas as pd # type: ignore
from django.db.models import Count, Q
from .models import *

def treatment_strategies(request):
    # Prepare the data for each treatment type
    treatment_types = ['Radiation Therapy', 'Chemotherapy', 'Hormone Therapy']
    treatment_data = {}

    # Get the count of treatment types from TreatmentTimeline
    for treatment_type in treatment_types:
        treatment_data[treatment_type] = TreatmentTimeline.objects.filter(treatment_type=treatment_type).distinct().count()

    # Prepare the pie charts
    pie_charts = []
    total_treatments = sum(treatment_data.values())  # Total of all treatments for reference

    for treatment_type in treatment_types:
        values = [treatment_data[treatment_type], total_treatments - treatment_data[treatment_type]]
        labels = ["Yes", "No"]  # Replace with 'Yes' and 'No'

        trace = go.Pie(
            labels=labels,
            values=values,
            hole=0.3,  # Donut-style pie chart
            marker=dict(colors=['#FF7F0E', '#1F77B4']),  # Custom colors for Yes and No
            hoverinfo='none',  # Disable default hoverinfo to use hovertemplate
            hovertemplate="<b>%{label}</b><br>Patients: %{value}<br>Percentage: %{percent:.1%}<extra></extra>",
            textinfo='label+percent'
        )

        layout = go.Layout(
            title=f"{treatment_type} Usage",
            margin=dict(l=10, r=10, b=20, t=50),
            showlegend=False,
        )

        fig = go.Figure(data=[trace], layout=layout)
        pie_charts.append(to_html(fig, full_html=False))

    # Render the pie charts in the template
    return render(request, 'treatment_strategies.html', {'pie_charts': pie_charts})

def survival_analysis(request):
    # Prepare lists for the terms and respective fields
    survival_terms = [
        {'status': 'os_status', 'months': 'os_months', 'title': 'Overall Survival (OS)'},
        {'status': 'dss_status', 'months': 'dss_months', 'title': 'Disease-Specific Survival (DSS)'},
        {'status': 'dfs_status', 'months': 'dfs_months', 'title': 'Disease-Free Survival (DFS)'},
        {'status': 'pfs_status', 'months': 'pfs_months', 'title': 'Progression-Free Survival (PFS)'}
    ]

    # Container for all the plots
    plots_html = []

    for term in survival_terms:
        # Fetch data for each term
        patients = ClinicalPatient.objects.all().values(term['status'], term['months'])

        # Convert to DataFrame
        data = pd.DataFrame(list(patients))
        if data[term['months']].isnull().all():
            # Skip empty data sets
            continue

        # Prepare the data for Kaplan-Meier fitting
        data = data.dropna(subset=[term['months'], term['status']])
        data['event'] = data[term['status']].apply(lambda x: 1 if '1' in str(x) else 0)
        data['duration'] = data[term['months']]

        # Fit the Kaplan-Meier model
        kmf = KaplanMeierFitter()
        kmf.fit(durations=data['duration'], event_observed=data['event'])

        # Prepare the plot data and labels
        survival_percentages = kmf.survival_function_['KM_estimate'] * 100
        at_risk_counts = kmf.event_table['at_risk']

        trace = go.Scatter(
            x=kmf.survival_function_.index,
            y=survival_percentages,
            mode='lines+markers',
            marker=dict(size=6, color='blue'),
            hoverinfo='text',
            text=[f"Time: {time} months<br>Survival Probability: {prob:.2f}%<br>Patients at risk: {at_risk}"
                  for time, prob, at_risk in zip(kmf.survival_function_.index, survival_percentages, at_risk_counts)]
        )

        layout = go.Layout(
            title=term['title'],
            xaxis=dict(title="Time (months)"),
            yaxis=dict(title="Survival Probability (%)"),
            hovermode="closest"
        )

        fig = go.Figure(data=[trace], layout=layout)
        plots_html.append(to_html(fig, full_html=False))

    # Pass the list of plots to the template
    return render(request, 'survival_analysis.html', {'plots_html': plots_html})

from django.shortcuts import render
from .models import ClinicalPatient
from django.db.models import Count, Case, When, Value

def clinical_dashboard(request):
    # Data for KM plot, converting None subtype values to "Unknown"
    km_data = ClinicalPatient.objects.values_list('os_months', 'os_status', 'subtype')
    km_data = [[months, status, subtype if subtype is not None else "Unknown"] for months, status, subtype in km_data]

    # Data for the subtype pie chart, handling None as "Unknown"
    subtype_data = ClinicalPatient.objects.annotate(
        subtype_label=Case(
            When(subtype__isnull=True, then=Value('Unknown')),
            default='subtype'
        )
    ).values('subtype_label').annotate(count=Count('subtype_label'))

    context = {
        'km_data': km_data,
        'subtype_data': list(subtype_data),
    }

    print(context['km_data'])
    return render(request, 'clinical_dashboard.html', context)