from django.db import models

class Dataset(models.Model):
    name = models.CharField(max_length=100, unique=True, help_text="Unique name for the dataset")
    description = models.TextField(null=True, blank=True, help_text="Description of the dataset")
    case_ids = models.ManyToManyField('Patient', related_name='datasets', help_text="Patients linked to this dataset")
    case_id_suffix = models.CharField(max_length=50, null=True, blank=True, help_text="Suffix for case IDs to distinguish this dataset")

    class Meta:
        db_table = 'tcga_dataset'
        verbose_name = 'Dataset'
        verbose_name_plural = 'Datasets'

    def __str__(self):
        return self.name

class Patient(models.Model):
    """
    Core patient model that serves as the main reference point for all related data.
    This model stores the unique patient identifier from TCGA dataset.
    """
    dataset = models.ForeignKey(Dataset,blank=True, null=True,default=None, on_delete=models.CASCADE, related_name='patients') # added this new line
    patient_id = models.CharField(max_length=50, primary_key=True, unique=True, 
                                help_text="Unique TCGA patient identifier")
    
    case_id = models.CharField(max_length=50, null=True, blank=True, 
                               help_text="Case ID associated with the patient (if applicable)")

    
    class Meta:
        db_table = 'tcga_patient'
        verbose_name = 'Patient'
        verbose_name_plural = 'Patients'
    
    def __str__(self):
        return self.patient_id

class ClinicalSample(models.Model):
    """
    Stores clinical sample data linked to each patient.
    """
    patient = models.ForeignKey(Patient, on_delete=models.CASCADE, related_name='clinical_samples')
    sample_id = models.CharField(max_length=50, unique=True)
    oncotree_code = models.CharField(max_length=20, null=True, blank=True)
    cancer_type = models.CharField(max_length=100, null=True, blank=True)
    cancer_type_detailed = models.CharField(max_length=100, null=True, blank=True)
    tumor_type = models.CharField(max_length=100, null=True, blank=True)
    grade = models.CharField(max_length=50, null=True, blank=True)
    tumor_tissue_site = models.CharField(max_length=100, null=True, blank=True)
    aneuploidy_score = models.FloatField(null=True, blank=True)
    sample_type = models.CharField(max_length=50, null=True, blank=True)
    msi_score_mantis = models.FloatField(null=True, blank=True)
    msi_sensor_score = models.FloatField(null=True, blank=True)
    somatic_status = models.CharField(max_length=50, null=True, blank=True)
    tmb_nonsynonymous = models.FloatField(null=True, blank=True)
    tissue_source_site = models.CharField(max_length=100, null=True, blank=True)

    class Meta:
        db_table = 'tcga_clinical_sample'
        verbose_name = 'Clinical Sample'
        verbose_name_plural = 'Clinical Samples'

    def __str__(self):
        return f"{self.patient_id} - {self.sample_id}"

class TreatmentTimeline(models.Model):
    """
    Stores treatment timeline data for each patient, including details about
    various types of treatments (radiation, chemotherapy, etc.)
    """
    TREATMENT_TYPES = [
        ('Radiation Therapy', 'Radiation Therapy'),
        ('Chemotherapy', 'Chemotherapy'),
        ('Hormone Therapy', 'Hormone Therapy'),
        ('Immunotherapy', 'Immunotherapy'),
        ('Targeted Molecular Therapy', 'Targeted Molecular Therapy'),
        ('Other', 'Other')
    ]

    EVENT_TYPES = [
        ('Treatment', 'Treatment'),
        ('Follow Up', 'Follow Up')
    ]

    patient = models.ForeignKey(Patient, on_delete=models.CASCADE, related_name='treatments')
    start_date = models.IntegerField(help_text="Days from initial diagnosis")
    stop_date = models.IntegerField(null=True, blank=True, help_text="Days from initial diagnosis")
    event_type = models.CharField(max_length=50, choices=EVENT_TYPES)
    treatment_type = models.CharField(max_length=50, choices=TREATMENT_TYPES)
    treatment_subtype = models.CharField(max_length=100, null=True, blank=True)
    agent = models.CharField(max_length=100, null=True, blank=True)
    number_of_cycles = models.CharField(max_length=50, null=True, blank=True)
    prescribed_dose = models.CharField(max_length=50, null=True, blank=True)
    prescribed_dose_units = models.CharField(max_length=20, null=True, blank=True)
    regimen_number = models.CharField(max_length=20, null=True, blank=True)
    regimen_indication = models.CharField(max_length=50, null=True, blank=True)
    measure_of_response = models.CharField(max_length=50, null=True, blank=True)
    clinical_trial_drug_classification = models.CharField(max_length=100, null=True, blank=True)
    route_of_administration = models.CharField(max_length=20, null=True, blank=True)
    therapy_ongoing = models.CharField(max_length=5, null=True, blank=True)
    total_dose = models.CharField(max_length=20, null=True, blank=True)
    total_dose_units = models.CharField(max_length=20, null=True, blank=True)
    tx_on_clinical_trial = models.CharField(max_length=5, null=True, blank=True)
    anatomic_treatment_site = models.CharField(max_length=100, null=True, blank=True)
    course_number = models.FloatField(null=True, blank=True)
    number_of_fractions = models.IntegerField(null=True, blank=True)
    radiation_dosage = models.IntegerField(null=True, blank=True)
    radiation_treatment_ongoing = models.CharField(max_length=5, null=True, blank=True)
    radiation_type = models.CharField(max_length=50, null=True, blank=True)
    radiation_type_notes = models.TextField(null=True, blank=True)
    radiation_units = models.CharField(max_length=20, null=True, blank=True)

    class Meta:
        db_table = 'tcga_treatment_timeline'
        verbose_name = 'Treatment Timeline'
        verbose_name_plural = 'Treatment Timelines'
        ordering = ['patient', 'start_date']

    def __str__(self):
        return f"{self.patient_id} - {self.treatment_type} ({self.start_date})"

class ClinicalPatient(models.Model):
    """
    Stores detailed clinical patient data including demographics, diagnosis, 
    and survival information.
    """
    LIVING_STATUS_CHOICES = [
        ('0:LIVING', 'Living'),
        ('1:DECEASED', 'Deceased')
    ]
    
    DISEASE_STATUS_CHOICES = [
        ('0:DiseaseFree', 'Disease Free'),
        ('1:Recurred/Progressed', 'Recurred/Progressed')
    ]
    
    PROGRESSION_STATUS_CHOICES = [
        ('0:CENSORED', 'Censored'),
        ('1:PROGRESSION', 'Progression')
    ]

    patient = models.OneToOneField(Patient, on_delete=models.CASCADE, related_name='clinical_data')
    subtype = models.CharField(max_length=50, null=True, blank=True)
    cancer_type_acronym = models.CharField(max_length=10, null=True, blank=True)
    age = models.IntegerField(null=True, blank=True)
    sex = models.CharField(max_length=10)
    ajcc_pathologic_tumor_stage = models.CharField(max_length=20, null=True, blank=True)
    ajcc_staging_edition = models.CharField(max_length=20, null=True, blank=True)
    days_last_followup = models.IntegerField(null=True, blank=True)
    days_to_birth = models.IntegerField(null=True, blank=True)
    ethnicity = models.CharField(max_length=50, null=True, blank=True)
    
    # Disease Classification
    icd_10 = models.CharField(max_length=10, null=True, blank=True)
    icd_o_3_histology = models.CharField(max_length=10, null=True, blank=True)
    icd_o_3_site = models.CharField(max_length=10, null=True, blank=True)
    
    # Disease Status
    new_tumor_event_after_initial_treatment = models.CharField(max_length=3, null=True, blank=True)
    path_m_stage = models.CharField(max_length=10, null=True, blank=True)
    path_n_stage = models.CharField(max_length=10, null=True, blank=True)
    path_t_stage = models.CharField(max_length=10, null=True, blank=True)
    person_neoplasm_cancer_status = models.CharField(max_length=20, null=True, blank=True)
    
    # Demographics
    race = models.CharField(max_length=50, null=True, blank=True)
    weight = models.FloatField(null=True, blank=True)
    genetic_ancestry_label = models.CharField(max_length=20, null=True, blank=True)
    
    # Survival Analysis
    os_status = models.CharField(max_length=20, choices=LIVING_STATUS_CHOICES, null=True, blank=True)
    os_months = models.FloatField(null=True, blank=True)
    dss_status = models.CharField(max_length=50, null=True, blank=True)
    dss_months = models.FloatField(null=True, blank=True)
    dfs_status = models.CharField(max_length=50, choices=DISEASE_STATUS_CHOICES, null=True, blank=True)
    dfs_months = models.FloatField(null=True, blank=True)
    pfs_status = models.CharField(max_length=50, choices=PROGRESSION_STATUS_CHOICES, null=True, blank=True)
    pfs_months = models.FloatField(null=True, blank=True)

    class Meta:
        db_table = 'tcga_clinical_patient'
        verbose_name = 'Clinical Patient Data'
        verbose_name_plural = 'Clinical Patient Data'
        indexes = [
            models.Index(fields=['subtype']),
            models.Index(fields=['os_status']),
            models.Index(fields=['ajcc_pathologic_tumor_stage'])
        ]

    def __str__(self):
        return f"{self.patient_id} - {self.subtype or 'Unknown Subtype'}"

class HypoxiaScore(models.Model):
    """
    Stores mRNA-based hypoxia scores for each patient using different scoring methods
    (Buffa, Winter, and Ragnum scores)
    """
    patient = models.OneToOneField(Patient, on_delete=models.CASCADE, related_name='hypoxia_scores')
    buffa_score = models.IntegerField(help_text="mRNA based Buffa Hypoxia Score")
    winter_score = models.IntegerField(help_text="mRNA based Winter Hypoxia Score")
    ragnum_score = models.IntegerField(help_text="mRNA based Ragnum Hypoxia Score")

    class Meta:
        db_table = 'tcga_hypoxia_scores'
        verbose_name = 'Hypoxia Score'
        verbose_name_plural = 'Hypoxia Scores'
        indexes = [
            models.Index(fields=['buffa_score']),
            models.Index(fields=['winter_score']),
            models.Index(fields=['ragnum_score'])
        ]

    def __str__(self):
        return f"{self.patient_id} - Buffa: {self.buffa_score}, Winter: {self.winter_score}, Ragnum: {self.ragnum_score}"

    @property
    def average_score(self):
        """Returns the average of all three hypoxia scores"""
        return (self.buffa_score + self.winter_score + self.ragnum_score) / 3

    def get_hypoxia_status(self):
        """
        Returns a classification based on average hypoxia score
        High: > 10
        Medium: -10 to 10
        Low: < -10
        """
        avg = self.average_score
        if avg > 10:
            return "High"
        elif avg < -10:
            return "Low"
        return "Medium"

class Gene(models.Model):
    """
    Represents a gene with its metadata.
    """
    hugo_symbol = models.CharField(max_length=100, null=True, blank=True, help_text="Gene name")
    entrez_gene_id = models.IntegerField(unique=True, help_text="Unique Entrez Gene ID", null=True, blank=True)
    
    class Meta:
        db_table = 'tcga_gene'
        verbose_name = 'Gene'
        verbose_name_plural = 'Genes'
    
    def __str__(self):
        return f"{self.hugo_symbol or 'Unnamed'} (ID: {self.entrez_gene_id})"

class GeneExpression(models.Model):
    """
    Represents gene expression data (z-scores) for a specific patient and gene.
    """
    gene = models.ForeignKey(Gene, on_delete=models.CASCADE, related_name="expressions")
    patient = models.ForeignKey('Patient', on_delete=models.CASCADE, related_name="gene_expressions")
    z_score = models.FloatField(help_text="Z-score for the gene expression")

    class Meta:
        verbose_name = 'Gene Expression'
        verbose_name_plural = 'Gene Expressions'
        unique_together = ('gene', 'patient')
    
    def __str__(self):
        return f"{self.gene.hugo_symbol or 'Unnamed'} - {self.patient.patient_id}: {self.z_score}"

class Mutation(models.Model):
    hugo_symbol = models.CharField(max_length=255)  # Gene name
    chromosome = models.CharField(max_length=5)  # Chromosome number
    start_position = models.PositiveIntegerField()  # Start position of mutation
    end_position = models.PositiveIntegerField()  # End position of mutation
    variant_classification = models.CharField(max_length=255)  # e.g., Missense_Mutation
    variant_type = models.CharField(max_length=50)  # e.g., SNP, INS, DEL
    reference_allele = models.CharField(max_length=50, blank=True, null=True)  # Reference base
    tumor_seq_allele1 = models.CharField(max_length=50, blank=True, null=True)  # Tumor allele 1
    tumor_seq_allele2 = models.CharField(max_length=50, blank=True, null=True)  # Tumor allele 2
    tumor_sample_barcode = models.ForeignKey(
        Patient, 
        on_delete=models.CASCADE, 
        related_name="mutations", 
        help_text="Reference to the patient associated with this mutation"
    )

    class Meta:
        verbose_name = "Mutation"
        verbose_name_plural = "Mutations"
        indexes = [
            models.Index(fields=['hugo_symbol']),
            models.Index(fields=['chromosome']),
        ]

    def __str__(self):
        return f"{self.hugo_symbol} ({self.chromosome}:{self.start_position}-{self.end_position})"

class Cna_hg19(models.Model):
    patient = models.ForeignKey(
        Patient, 
        on_delete=models.CASCADE, 
        related_name="cna_hg19", 
        help_text="Reference to the patient associated with this"
    )
    chromosome = models.CharField(max_length=5)
    loc_start = models.PositiveIntegerField()
    loc_end = models.PositiveIntegerField()
    num_mark = models.FloatField()
    seg_mean = models.FloatField()

    class Meta:
        verbose_name = 'CNA HG19'
        verbose_name_plural = 'CNA HG19'
        indexes = [
            models.Index(fields=['patient_id']),
            # models.Index(fields=['chromosome']),
        ]

    def __str__(self):
        return f"{self.patient.patient_id} - {self.chromosome}:{self.loc_start}-{self.loc_end}"

class SvData(models.Model):
    patient = models.ForeignKey(
        Patient, 
        on_delete=models.CASCADE, 
        related_name="sv_data", 
        help_text="Reference to the patient associated with this"
    )
    site1_gene = models.ForeignKey(
        Gene, 
        on_delete=models.CASCADE, 
        related_name="site1_sv_data", 
        help_text="Reference to the gene at site 1"
    )
    site2_gene = models.ForeignKey(
        Gene, 
        on_delete=models.CASCADE, 
        related_name="site2_sv_data", 
        help_text="Reference to the gene at site 2"
    )
    sv_status = models.CharField(max_length=50, help_text="Status of the structural variant (e.g., SOMATIC)")

    class Meta:
        verbose_name = 'Structural Variant Data'
        verbose_name_plural = 'Structural Variant Data'
        indexes = [
            models.Index(fields=['patient_id']),
            # models.Index(fields=['site1_gene']),
            # models.Index(fields=['site2_gene']),
        ]

    def __str__(self):
        return f"{self.patient} - {self.site1_gene.hugo_symbol} & {self.site2_gene.hugo_symbol} ({self.sv_status})"
    
class Cna(models.Model):
    gene = models.ForeignKey(Gene, on_delete=models.CASCADE, related_name="cna")
    patient = models.ForeignKey('Patient', on_delete=models.CASCADE, related_name="cna")
    value = models.IntegerField(help_text="cna_value for the gene expression")

    class Meta:
        verbose_name = 'CNA'
        verbose_name_plural = 'CNA'
        unique_together = ('gene', 'patient')
    
    def __str__(self):
        return f"{self.gene.hugo_symbol or 'Unnamed'} - {self.patient.patient_id}: {self.value}"