# application/admin.py

from django.contrib import admin
from django.db.models import Avg
from .models import Cna, Cna_hg19, Mutation, Patient, ClinicalPatient, ClinicalSample, SvData, TreatmentTimeline, HypoxiaScore, Gene, GeneExpression,Dataset

admin.site.register(Dataset)
admin.site.register(Gene)
admin.site.register(GeneExpression)
admin.site.register(Mutation)
admin.site.register(SvData)
admin.site.register(Cna_hg19)
admin.site.register(Cna)

class HypoxiaStatusFilter(admin.SimpleListFilter):
    title = 'Hypoxia Status'
    parameter_name = 'hypoxia_status'

    def lookups(self, request, model_admin):
        return (
            ('high', 'High'),
            ('medium', 'Medium'),
            ('low', 'Low'),
        )

    def queryset(self, request, queryset):
        if self.value() == 'high':
            return queryset.filter(buffa_score__gt=10)
        elif self.value() == 'low':
            return queryset.filter(buffa_score__lt=-10)
        elif self.value() == 'medium':
            return queryset.filter(buffa_score__gte=-10, buffa_score__lte=10)


@admin.register(Patient)
class PatientAdmin(admin.ModelAdmin):
    list_display = ('patient_id',)
    search_fields = ('patient_id',)


@admin.register(ClinicalPatient)
class ClinicalPatientAdmin(admin.ModelAdmin):
    list_display = (
        'patient_id', 'age', 'sex', 'subtype', 
        'ajcc_pathologic_tumor_stage', 'os_status', 
        'os_months'
    )
    list_filter = (
        'sex', 
        'subtype', 
        'ajcc_pathologic_tumor_stage',
        'os_status',
        'race',
        'ethnicity'
    )
    search_fields = ('patient_id', 'subtype')
    readonly_fields = ('patient',)
    
    fieldsets = (
        ('Patient Information', {
            'fields': ('patient', 'subtype', 'age', 'sex', 'race', 'ethnicity', 'weight')
        }),
        ('Disease Classification', {
            'fields': (
                'cancer_type_acronym',
                'ajcc_pathologic_tumor_stage',
                'ajcc_staging_edition',
                'icd_10',
                'icd_o_3_histology',
                'icd_o_3_site'
            )
        }),
        ('Disease Status', {
            'fields': (
                'path_m_stage',
                'path_n_stage',
                'path_t_stage',
                'person_neoplasm_cancer_status',
                'new_tumor_event_after_initial_treatment'
            )
        }),
        ('Survival Data', {
            'fields': (
                'os_status',
                'os_months',
                'dss_status',
                'dss_months',
                'dfs_status',
                'dfs_months',
                'pfs_status',
                'pfs_months'
            )
        })
    )


@admin.register(ClinicalSample)
class ClinicalSampleAdmin(admin.ModelAdmin):
    list_display = (
        'patient_id', 'sample_id', 'cancer_type',
        'tumor_type', 'sample_type'
    )
    list_filter = (
        'cancer_type',
        'tumor_type',
        'sample_type'
    )
    search_fields = (
        'patient__patient_id',
        'sample_id',
        'cancer_type'
    )
    readonly_fields = ('patient',)


@admin.register(TreatmentTimeline)
class TreatmentTimelineAdmin(admin.ModelAdmin):
    list_display = (
        'patient_id', 'treatment_type', 'start_date',
        'stop_date', 'therapy_ongoing'
    )
    list_filter = (
        'treatment_type',
        'event_type',
        'therapy_ongoing',
        'radiation_treatment_ongoing'
    )
    search_fields = (
        'patient__patient_id',
        'treatment_type',
        'agent'
    )
    readonly_fields = ('patient',)
    
    fieldsets = (
        ('Basic Information', {
            'fields': ('patient', 'event_type', 'treatment_type', 'treatment_subtype')
        }),
        ('Timeline', {
            'fields': ('start_date', 'stop_date')
        }),
        ('Treatment Details', {
            'fields': (
                'agent',
                'number_of_cycles',
                'prescribed_dose',
                'prescribed_dose_units',
                'total_dose',
                'total_dose_units',
                'therapy_ongoing'
            )
        }),
        ('Radiation Specific', {
            'fields': (
                'radiation_dosage',
                'radiation_type',
                'radiation_type_notes',
                'radiation_units',
                'radiation_treatment_ongoing',
                'anatomic_treatment_site',
                'number_of_fractions'
            )
        })
    )


@admin.register(HypoxiaScore)
class HypoxiaScoreAdmin(admin.ModelAdmin):
    list_display = (
        'patient_id', 'buffa_score', 'winter_score',
        'ragnum_score', 'get_average_score', 'get_hypoxia_status'
    )
    list_filter = (HypoxiaStatusFilter,)
    search_fields = ('patient__patient_id',)
    readonly_fields = ('patient', 'get_average_score', 'get_hypoxia_status')

    def get_average_score(self, obj):
        return round((obj.buffa_score + obj.winter_score + obj.ragnum_score) / 3, 2)
    get_average_score.short_description = 'Average Score'
    
    fieldsets = (
        ('Patient', {
            'fields': ('patient',)
        }),
        ('Hypoxia Scores', {
            'fields': (
                'buffa_score',
                'winter_score',
                'ragnum_score'
            )
        }),
        ('Calculated Values', {
            'fields': (
                'get_average_score',
                'get_hypoxia_status'
            )
        })
    )


# Customize admin site header and title
admin.site.site_header = 'TCGA Breast Cancer Database Administration'
admin.site.site_title = 'TCGA Admin Portal'
admin.site.index_title = 'Database Management'