import pandas as pd
import numpy as np
import logging
import re 
from typing import Optional
from .general import DataProcessorGeneral

logging.basicConfig(
    level = logging.INFO,                                 
    format = '%(asctime)s - %(levelname)s - %(message)s'  
)

class DataProcessorHeadNeck(DataProcessorGeneral):
    def __init__(self):
        super().__init__() 

        # head neck-specific attributes
        self.enhanced_df = None

    def process_enhanced(self,
                         file_path: str,
                         patient_ids: list = None,
                         drop_stages: bool = True, 
                         drop_surgery_type: bool = True,
                         drop_dates: bool = True) -> Optional[pd.DataFrame]: 
        """
        Processes Enhanced_AdvUrothelial.csv to standardize categories, consolidate 
        staging information, and calculate time-based metrics between key clinical events.

        Parameters
        ----------
        file_path : str
            Path to Enhanced_AdvUrothelial.csv file
        patient_ids : list, optional
            List of PatientIDs to process. If None, processes all patients
        drop_stages : bool, default=True
            If True, drops original staging columns (GroupStage, TStage, NStage, and MStage) after creating modified versions
        drop_surgery_type : bool, default=True
            If True, drops original surgery type after creating modified version
        drop_dates : bool, default=True
            If True, drops date columns (DiagnosisDate, AdvancedDiagnosisDate, and SurgeryDate) after calculating durations

        Returns
        -------
        pd.DataFrame or None
            - PatientID : object
                unique patient identifier
            - PrimarySite : category
                anatomical site of cancer
            - SmokingStatus : category
                smoking history
            - Surgery : Int64
                binary indicator (0/1) for whether patient had surgery
            - SurgeryType_mod : category
                consolidated surgery type
            - DiseaseGrade : category
                tumor grade (high, low, and unknown) at time of first diagnosis
            - GroupStage_mod : category
                consolidated overall staging (0-IV and unknown) at time of first diagnosis
            - TStage_mod : category
                consolidated tumor staging (T1-T4, non-muscle invasive bladder cancer [nmibc], other) at time of first diagnosis
            - NStage_mod : category
                consolidated lymph node staging (N0, N1, N2, N3, and unknown) at time of first diagnosis
            - MStage_mod : category
                consolidated metastasis staging (M0, M1, and unknown) at time of first diagnosis
            - days_diagnosis_to_adv : float
                days from first diagnosis to advanced disease 
            - adv_diagnosis_year : category
                year of advanced diagnosis 
            - days_diagnosis_to_surgery : float
                days from first diagnosis to surgery 
            
            Original staging and date columns retained if respective drop_* = False

        Notes
        -----
        Notable T-Stage consolidation decisions:
            - Ta and Tis : 'nmibc' (Non-muscle invasive disease)
            - T0, TX, Unknown/not documented : 'other' 

        If days_diagnosis_to_surgery is < 0, the value is set to 0; otherwise, the value is left as is. 
        Cases where days_diagnosis_to_surgery is < 0 occur when the date of surgery is set to the first of the year 
        and likely reflects uknown exact date of surgery. By setting these cases to 0, we'll presume that the 
        surgery date occurred on the diagnosis date. 

        Output handling:
        - Duplicate PatientIDs are logged as warnings if found but reatained in output
        - Processed DataFrame is stored in self.enhanced_df
        """
        # Input validation
        if patient_ids is not None:
            if not isinstance(patient_ids, list):
                raise TypeError("patient_ids must be a list or None")
            
        try:
            df = pd.read_csv(file_path)
            logging.info(f"Successfully read Enhanced_AdvUrothelial.csv file with shape: {df.shape} and unique PatientIDs: {(df['PatientID'].nunique())}")

            # Filter for specific PatientIDs if provided
            if patient_ids is not None:
                logging.info(f"Filtering for {len(patient_ids)} specific PatientIDs")
                df = df[df['PatientID'].isin(patient_ids)]
                logging.info(f"Successfully filtered Enhanced_AdvUrothelial.csv file with shape: {df.shape} and unique PatientIDs: {(df['PatientID'].nunique())}")
        
            # Convert categorical columns
            categorical_cols = ['PrimarySite', 
                                'DiseaseGrade',
                                'GroupStage',
                                'TStage', 
                                'NStage',
                                'MStage', 
                                'SmokingStatus', 
                                'SurgeryType']
        
            df[categorical_cols] = df[categorical_cols].astype('category')

            # Recode stage variables using class-level mapping and create new column
            df['GroupStage_mod'] = df['GroupStage'].map(self.GROUP_STAGE_MAPPING).astype('category')
            df['TStage_mod'] = df['TStage'].map(self.T_STAGE_MAPPING).astype('category')
            df['NStage_mod'] = df['NStage'].map(self.N_STAGE_MAPPING).astype('category')
            df['MStage_mod'] = df['MStage'].map(self.M_STAGE_MAPPING).astype('category')

            # Drop original stage variables if specified
            if drop_stages:
                df = df.drop(columns=['GroupStage', 'TStage', 'NStage', 'MStage'])

            # Recode surgery type variable using class-level mapping and create new column
            df['SurgeryType_mod'] = df['SurgeryType'].map(self.SURGERY_TYPE_MAPPING).astype('category')

            # Drop original surgery type variable if specified
            if drop_surgery_type:
                df = df.drop(columns=['SurgeryType'])

            # Convert date columns
            date_cols = ['DiagnosisDate', 'AdvancedDiagnosisDate', 'SurgeryDate']
            for col in date_cols:
                df[col] = pd.to_datetime(df[col])
            
            # Convert boolean column to binary (0/1)
            df['Surgery'] = df['Surgery'].astype('Int64')

            # Generate new variables 
            df['days_diagnosis_to_adv'] = (df['AdvancedDiagnosisDate'] - df['DiagnosisDate']).dt.days
            df['adv_diagnosis_year'] = pd.Categorical(df['AdvancedDiagnosisDate'].dt.year)
            df['days_diagnosis_to_surgery'] = (df['SurgeryDate'] - df['DiagnosisDate']).dt.days
            
            # For those with value <0, set to 0; otherwise, leave as is
            df['days_diagnosis_to_surgery'] = np.where(df['days_diagnosis_to_surgery'] < 0, 0, df['days_diagnosis_to_surgery'])
    
            if drop_dates:
                df = df.drop(columns = ['AdvancedDiagnosisDate', 'DiagnosisDate', 'SurgeryDate'])

            # Check for duplicate PatientIDs
            if len(df) > df['PatientID'].nunique():
                duplicate_ids = df[df.duplicated(subset = ['PatientID'], keep = False)]['PatientID'].unique()
                logging.warning(f"Duplicate PatientIDs found: {duplicate_ids}")

            logging.info(f"Successfully processed Enhanced_AdvUrothelial.csv file with final shape: {df.shape} and unique PatientIDs: {(df['PatientID'].nunique())}")
            self.enhanced_df = df
            return df

        except Exception as e:
            logging.error(f"Error processing Enhanced_AdvUrothelial.csv file: {e}")
            return None