import os
import logging
import re
import pandas as pd
import numpy as np
from joblib import Parallel, delayed
from typing import List, Optional
from plinkformatter.plink_utils import (
    generate_bed_bim_fam,
    calculate_kinship_matrix
)
from plinkformatter.utils import validate_files_exist_and_not_empty
from plinkformatter.exceptions import PLINK2Error, PhenotypeMeasureError


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


def generate_map(df: pd.DataFrame, map_file: str) -> None:
    """
    Generates a PLINK-formatted .map file from the provided DataFrame.

    Args:
        df: DataFrame containing genotype information with at least the
            columns ['chr', 'rs', 'bp38'].
        map_file: Path to the output MAP file.

    Returns:
        None. The function writes the .map file to the specified location.
    """
    # Ensure missing 'rs' values are replaced with a combination of 'chr' and 'bp38'
    df['rs'] = df.apply(lambda row: f"{row['chr']}_{row['bp38']}" if pd.isna(row['rs']) or row['rs'] == '' else row['rs'], axis=1)
    
    # Ensure no duplicates based on 'chr', 'rs', 'bp38'
    df = df.drop_duplicates(subset=['chr', 'rs', 'bp38'])

    # Prepare the map DataFrame with necessary columns and add cM column
    map_df = df[['chr', 'rs', 'bp38']].copy()
    map_df.insert(2, 'cM', 0)

    # Write to file
    map_df.to_csv(map_file, sep="\t", index=False, header=False)


def apply_transformation(df: pd.DataFrame, reference_col: str) -> pd.DataFrame:
    """
    Applies genotype transformation based on a reference column.

    Transformation rules:
    - If a cell matches the reference column value, set '1 1'.
    - If a cell is 'N', set '0 0'.
    - If a cell is 'H', set '1 2'.
    - Otherwise, set '2 2'.

    Args:
        df: DataFrame containing genotype data to transform.
        reference_col: The column name in df to be used as a reference for
            comparison.

    Returns:
        DataFrame: The transformed DataFrame.
    """
    df_transformed = df.copy()

    # Get the reference column values for comparison
    reference_values = df[reference_col]
 
    # Apply the condition where the value matches the reference column
    for col in df.columns:
        # If values match the reference column (b_6), set '1 1'
        df_transformed.loc[df[col] == reference_values, col] = "1 1"
        # Set '0 0' for 'N' values
        df_transformed.loc[df[col] == "N", col] = "0 0"
        # Set '1 2' for 'H' values
        df_transformed.loc[df[col] == "H", col] = "1 2"
        # For all other values, set '2 2'
        df_transformed.loc[(df[col] != reference_values) & (df[col] != "N") & (df[col] != "H"), col] = "2 2"

    return df_transformed


def generate_ped(df: pd.DataFrame, ped_file: str) -> None:
    """
    Generates a PLINK-formatted .ped file from the given DataFrame.

    Args:
        df: DataFrame containing genotype data with at least an 'observed' column.
        ped_file: Path to the output PED file.

    Returns:
        None. The function writes the .ped file to the specified location.
    """
    # Find observed column
    index = df.columns.get_loc("observed")
    ped_df = df.iloc[:, (index + 1):].copy()
    
    # Rename "C57BL/6J" to "b_6"
    ped_df.rename(columns={"C57BL/6J": "b_6"}, inplace=True)
    
    # Reference column for comparison
    ref_column = "b_6"

    # Apply genotype transformation row-wise for each column
    ped_df = apply_transformation(ped_df, ref_column)

    # Rename "b_6" back to "C57BL/6J"
    ped_df.rename(columns={"b_6": "C57BL/6J"}, inplace=True)
    
    # Apply transpose
    ped_df = ped_df.T

    # Add necessary columns for PLINK format
    ped_df.insert(0, "FID", ped_df.index)
    ped_df.insert(1, "IID", ped_df["FID"])
    ped_df.insert(2, "PID", 0)
    ped_df.insert(3, "MID", 0)
    ped_df.insert(4, "SEX", 0)
    ped_df.insert(5, "PHE", -9)

    # Write to file
    ped_df.to_csv(ped_file, sep="\t", index=False, header=False)


def generate_ped_map(
    genotype_df: pd.DataFrame,
    ped_file: str,
    map_file: str,
    gen_map: bool = True,
    gen_ped: bool = True
) -> None:
    """
    Generates both PLINK-formatted .ped and .map files from a genotype DataFrame.

    Args:
        genotype_df: DataFrame containing genotype data with columns ['chr', 'rs', 'bp38'].
        ped_file: Path to the output PED file.
        map_file: Path to the output MAP file.
        gen_map: Boolean indicating whether to generate the MAP file (default: True).
        gen_ped: Boolean indicating whether to generate the PED file (default: True).

    Returns:
        None. Writes the .ped and/or .map files to the specified locations.
    """
    # Filter DataFrame to exclude mitochondrial chromosomes
    df = genotype_df[genotype_df["chr"] != "M"].copy()
    
    # Ensure 'chr' and 'bp38' columns are strings for consistent processing
    df["chr"] = df["chr"].astype(str)
    df["bp38"] = df["bp38"].astype(str)
    
    # Replace missing or empty 'rs' values with concatenated 'chr' and 'bp38'
    df["rs"] = np.where(df["rs"] == "", df["chr"] + "_" + df["bp38"], df["rs"])
    
    # Drop duplicates based on 'chr', 'rs', and 'bp38' to ensure uniqueness
    df = df.drop_duplicates(subset=["chr", "rs", "bp38"])
    
    # Generate map file if requested
    if gen_map:
        generate_map(df, map_file)
    
    # Generate ped file if requested
    if gen_ped:
        generate_ped(df, ped_file)


def extract_pheno_measure(directory: str, measure_id: Optional[List[str]] = None) -> pd.DataFrame:
    """
    Extracts phenotype measures from local CSV files in the given directory.

    Args:
        directory: Directory containing CSV files for each measure.
        measure_id: Optional list of measure IDs to extract. If None, no measures will be extracted.

    Returns:
        pd.DataFrame: Concatenated phenotype data from all measure files, or an empty DataFrame if no files are found.
    """
    if measure_id is None:
        measure_id = []

    dat_list = []
    for mid in measure_id:
        try:
            file_path = os.path.join(directory, f"{mid}.csv")
            if os.path.exists(file_path):
                df = pd.read_csv(file_path)
                dat_list.append(df)
            else:
                print(f"File {file_path} does not exist.")
        except Exception as e:
            raise PhenotypeMeasureError(f"Error extracting phenotype for measure {mid}: {e}")

    if not dat_list:
        return pd.DataFrame()

    dat_all = pd.concat(dat_list, ignore_index=True)
    dat_all['strain'] = dat_all['strain'].str.replace(" ", "")
    dat_all['strain'] = dat_all['strain'].apply(lambda x: re.sub(r'.$', '', x) if x.startswith("CC") else x)

    return dat_all


def column_splitter(df: pd.DataFrame, split_index: int = 6) -> pd.DataFrame:
    """
    Splits columns in the DataFrame with the format "2 2" into two separate columns.

    Args:
        df: DataFrame with columns to be split.
        split_index: Index of the first column to split (default: 6).

    Returns:
        pd.DataFrame: DataFrame with split columns.
    """
    # Split the DataFrame into the initial columns and the columns that need splitting
    res = df.iloc[:, :split_index]
    sus = df.iloc[:, split_index:]

    # Create an empty list to store all new columns
    new_columns = []

    # Loop through each column in `sus` and split the values
    for x in sus.columns:
        split_cols = sus[x].str.split(" ", expand=True)
        # Append the split columns to the list
        new_columns.append(split_cols)

    # Concatenate the original columns and the new columns along axis 1
    new_res = pd.concat([res] + new_columns, axis=1)

    # Copy to de-fragment
    new_res = new_res.copy()

    new_res.columns = [i for i, _ in enumerate(new_res.columns)]

    return new_res


def generate_pheno_plink(
    ped_file: str,
    map_file: str,
    pheno: pd.DataFrame,
    outdir: str,
    ncore: int
) -> pd.DataFrame:
    """
    Generates PLINK-formatted .ped, .map, and .pheno files for phenotype data, stratified by sex and measurement number.

    Args:
        ped_file: Path to the input PED file.
        map_file: Path to the input MAP file.
        pheno: DataFrame containing phenotype data with columns ['strain', 'sex', 'measnum', 'zscore', 'value'].
        outdir: Directory where output files will be saved.
        ncore: Number of cores to use for parallel processing.

    Returns:
        pheno_s: Processed phenotype data filtered and sorted by strain.
    """
    # Ensure the output directory exists
    if not os.path.exists(outdir):
        os.makedirs(outdir)

    # Normalize strain names in phenotype data
    pheno['strain'] = pheno['strain'].str.replace(" ", "").str.upper()

    # Read map file
    map_df = pd.read_csv(map_file, header=None, sep="\t")
    map_df[1] = np.where(map_df[1] == ".", map_df[0].astype(str) + "_" + map_df[3].astype(str), map_df[1].astype(str))
    logging.debug("map file read and santized")
    
    # Read ped file
    ped_df = pd.read_csv(ped_file, header=None, sep='\t')
    ped_df[0] = ped_df[0].str.replace("?", "").str.replace(" ", "").str.upper()
    ped_df[1] = ped_df[1].str.replace("?", "").str.replace(" ", "").str.upper()
    logging.debug("ped file read and santized")

    # Filter and arrange pheno data by strain
    pheno_s = pheno[(pheno["sex"].isin(["f", "m"])) & (pheno["strain"].isin(ped_df[0]))]
    pheno_s = pheno_s.sort_values(by="strain")
    logging.debug("Filtered and sorted phenotype data by strain.")

    # Count number of each strain
    pheno_n = pheno_s["strain"].value_counts().reset_index()
    pheno_n.columns = ["strain", "count"]
    pheno_n = pheno_n.sort_values(by="strain")
    logging.debug("Counted number of occurrences for each strain.")

    # Subset ped data
    ped_s = ped_df[ped_df[0].isin(pheno_s["strain"])]
    ped_s = ped_s.sort_values(by=0).reset_index(drop=True)

    assert pheno_n["strain"].tolist() == ped_s[0].tolist(), "Strain mismatch between pheno and ped data after sorting."
    logging.debug("Verified that strains in phenotype and PED data match.")

    # Uncount ped data by number of each strain
    pheno_ped_n = ped_s.loc[ped_s.index.repeat(pheno_n["count"].values)].reset_index(drop=True)
    logging.debug("Expanded PED data according to strain counts.")

    # Combine pheno_s and pheno_ped_n
    pheno_ped = pd.concat([pheno_s.reset_index(drop=True), pheno_ped_n.reset_index(drop=True)], axis=1)
    pheno_ped[4] = np.where(pheno_ped["sex"] == "f", 2, 1)  # Set sex in ped file
    pheno_ped[5] = pheno_ped["zscore"]  # Set zscore as the phenotype column in ped file

    # Ensure FID and IID columns are present
    pheno_ped["FID"] = pheno_ped[0]
    pheno_ped["IID"] = pheno_ped[1]

    logging.debug("Combined phenotype data with PED data and set necessary columns.")

    # Write files for each measure
    def write_files(x):
        df = pheno_ped[pheno_ped["measnum"] == x]
        # Fix: remove iloc
        if "f" in df["sex"].unique():
            df_f = df[df["sex"] == "f"]

            # write pheno file
            df_f[["FID", "IID", "zscore", "value"]].to_csv(os.path.join(outdir, f"{x}.f.pheno"), sep=" ", index=False, header=False)

            # Fix:  I need only the integer indices
            indices = [x for x in df_f.columns if isinstance(x, (int))]
            df_f = df_f[indices]

            # Fix: need to split single columns into two columns: "2 2" -> 2 2
            df_f = column_splitter(df_f)

            # Write rest of files in PLINK format
            df_f.to_csv(os.path.join(outdir, f"{x}.f.ped"), sep=" ", index=False, header=False)
            map_df.to_csv(os.path.join(outdir, f"{x}.f.map"), sep="\t", index=False, header=False)
        
        if "m" in df["sex"].unique():
            df_m = df[df["sex"] == "m"]

            # write pheno file
            df_m[["FID", "IID", "zscore", "value"]].to_csv(os.path.join(outdir, f"{x}.m.pheno"), sep=" ", index=False, header=False)

            # Fix:  I need only the integer indices
            indices = [x for x in df_m.columns if isinstance(x, (int))]
            df_m = df_m[indices]
            
            # Fix: need to split single columns into two columns: "2 2" -> 2 2
            df_m = column_splitter(df_m)

            # Write rest of files in PLINK format
            df_m.to_csv(os.path.join(outdir, f"{x}.m.ped"), sep=" ", index=False, header=False)
            map_df.to_csv(os.path.join(outdir, f"{x}.m.map"), sep="\t", index=False, header=False)

    Parallel(n_jobs=ncore)(delayed(write_files)(x) for x in pheno_ped["measnum"].unique())

    return pheno_s


def prepare_pylmm_inputs(
    ped_file: str,
    map_file: str,
    measure_id_directory: str,
    measure_ids: List,
    outdir: str,
    ncore: int,
    plink2_path: str,
    genotype_df: Optional[pd.DataFrame] = None
) -> None:
    """
    Prepare the inputs required for running pylmm analysis.

    Args:
        ped_file: Path to the output PED file.
        map_file: Path to the output MAP file.
        measure_id_directory: Directory containing the measure csvs.
        measure_ids: List of measure ids.
        outdir: Output directory for generated files.
        ncore: Number of cores for parallel processing.
        plink2_path: Path to the PLINK2 executable.
        genotype_df [Optional]: DataFrame containing genotype data.  Can be supplied 
            in the rare cases we want to regenerate the ped/ map files.
    """
    # Ensure the output directory exists
    if not os.path.exists(outdir):
        os.makedirs(outdir)
    logging.debug("ensuring directory exists [complete]")
    # Default behavior is to assume ped/ map are already generated and reusable
    if genotype_df is not None:
        generate_ped_map(genotype_df, ped_file, map_file)

    # Step 1: Extract phenotype measures
    pheno = extract_pheno_measure(measure_id_directory, measure_ids)
    logging.debug("extract phenotype measures [complete]")
    # Step 2: Generate phenotype PLINK files
    generate_pheno_plink(ped_file, map_file, pheno, outdir, ncore)
    logging.debug("generate pheno plink [complete]")
    # Step 3: Generate BED/BIM/FAM files using PLINK2

    # iterate over all ,measure id files in directory
    for measure_id in measure_ids:
        for sex in ("f", "m"):
            measure_ped_file = os.path.join(outdir, f"{measure_id}.{sex}.ped")
            measure_map_file = os.path.join(outdir, f"{measure_id}.{sex}.map")
            output_prefix = os.path.join(outdir,  f"{measure_id}.{sex}")
            logging.debug("Creating PLINK for: %s.%s", measure_id, sex)

            if not os.path.exists(measure_ped_file):
                logging.debug(f"Data not available for sex: {sex}")
                break
            
            generate_bed_bim_fam(
                plink2_path,
                measure_ped_file,
                measure_map_file,
                output_prefix
            )

            logging.debug("bed bim fam [complete]")

            # Step 4: Calculate kinship matrix using PLINK2
            kin_prefix = os.path.join(outdir, f"{measure_id}.{sex}.kin")
            try:
                calculate_kinship_matrix(plink2_path, output_prefix, kin_prefix)
                logging.debug("kinship [complete]")

                # Ensure the kinship matrix files are created and not empty
                validate_files_exist_and_not_empty(kin_prefix, ['rel', 'rel.id'])
            except Exception as e:
                raise PLINK2Error(f"Error in PLINK2 processing for measure {measure_id}.{sex}: {e}")

    logging.debug("Directory contents: %s", os.listdir(outdir))


def extract_pheno_measure_from_api(url="https://phenome.jax.org/", measure_id=None):
    """
    Extract phenotype based on measure id.
    (Python version of Hao's R function.)
    
    :param url: url of MPD. Default: "https://phenome.jax.org/"
    :param measure_id: measure ids
    
    :return: One dataframe for all phenotype measures.
    """
    raise NotImplementedError("This function needs more testing in order to safely use it.")
    import requests
    from io import StringIO
    if measure_id is None:
        measure_id = []

    dat_list = []
    for mid in measure_id:
        try:
            response = requests.get(f"{url}api/pheno/animalvals/{mid}?csv=yes")
            if response.status_code == 200:
                csv_data = StringIO(response.text)
                df = pd.read_csv(csv_data)
                dat_list.append(df)
            else:
                print(f"Failed to fetch data for measure ID {mid}: Status code {response.status_code}")
        except Exception as e:
            print(f"Error fetching data for measure ID {mid}: {e}")

    if not dat_list:
        return pd.DataFrame()

    dat_all = pd.concat(dat_list, ignore_index=True)
    dat_all['strain'] = dat_all['strain'].str.replace(" ", "")
    dat_all['strain'] = dat_all['strain'].apply(lambda x: re.sub(r'.$', '', x) if x.startswith("CC") else x)

    return dat_all