import os
import csv
import logging
import operator

import h5py
import numpy as np
import pandas as pd
from typing.io import TextIO
from typing import List, Tuple

from ml4h.metrics import sparse_cross_entropy
from ml4h.TensorMap import TensorMap, Interpretation
from ml4h.DatabaseClient import BigQueryDatabaseClient, DatabaseClient
from ml4h.defines import TENSOR_MAPS_FILE_NAME, dataset_name_from_meaning
from ml4h.defines import DICTIONARY_TABLE, CODING_TABLE, PHENOTYPE_TABLE, JOIN_CHAR
from ml4h.tensormap.general import build_tensor_from_file, build_categorical_tensor_from_file
from ml4h.tensormap.text import random_text_window_tensor, token_dictionary_and_text_from_file, token_dictionary_from_hd5_key, random_array_window_tensors
from ml4h.tensorize.tensor_writer_ukbb import disease_prevalence_status, get_disease2tsv, disease_incidence_status, disease_censor_status


LESS_THAN_CODES = "('Less than a year', 'Less than once a week', 'Less than one mile', 'Less than an hour a day', 'Less than one a day', 'Less than one', 'Less than once a year', 'Less than 1 year ago', 'Less than a year ago', 'Less than one year', 'Less than one cigarette per day')"


def make_test_tensor_maps(desired_map_name: str) -> TensorMap:
    for n in range(1, 6):
        if desired_map_name == f'{n}d_cont':
            return TensorMap(f'{n}d_cont', shape=tuple(range(2, n + 2)), interpretation=Interpretation.CONTINUOUS)
        if desired_map_name == f'{n}d_cat':
            return TensorMap(
                f'{n}d_cat', shape=tuple(range(2, n + 2)), interpretation=Interpretation.CATEGORICAL,
                channel_map={f'c_{i}': i for i in range(n + 1)},
            )


def write_tensor_maps(args) -> None:
    logging.info("Making tensor maps...")

    tensor_maps_file = f"{args.output_folder}/{TENSOR_MAPS_FILE_NAME}.py"
    db_client = BigQueryDatabaseClient(credentials_file=args.bigquery_credentials_file)
    with open(tensor_maps_file, 'w') as f:
        f.write(_get_tensor_map_file_imports())
        #_write_disease_tensor_maps(args.phenos_folder, f)
        #_write_disease_tensor_maps_incident_prevalent(args.phenos_folder, f)
        #_write_phecode_tensor_maps(f, args.phecode_definitions, db_client)
        _write_continuous_tensor_maps(f, db_client)

        f.write('\n')
        logging.info(f"Wrote the tensor maps to {tensor_maps_file}.")


def _get_tensor_map_file_imports() -> str:
    return (
        f"# TensorMaps automatically generated by {os.path.basename(__file__)}\n"
        f"# DO NOT EDIT\n\n"
        f"from ml4h.defines import StorageType\n"
        f"from ml4h.metrics import weighted_crossentropy\n"
        f"from ml4h.tensormap.ukb.demographics import prevalent_incident_tensor\n"
        f"from ml4h.TensorMap import TensorMap, Interpretation, make_range_validator\n\n"
    )


def _write_disease_tensor_maps(phenos_folder: str, f: TextIO) -> None:
    f.write(f"\n\n#  TensorMaps for MPG disease phenotypes\n")
    disease2tsv = get_disease2tsv(phenos_folder)
    logging.info(f"Got:{len(disease2tsv)} disease TSVs from:{phenos_folder}")
    status = disease_censor_status(disease2tsv, 1000000, 5000000)
    logging.info(f"Got status for all diseases.")
    for d in sorted(list(disease2tsv.keys())):
        total = len(status[d])
        diseased = np.sum(list(status[d].values()))
        factor = int(total / (1 + diseased * 2))
        f.write(
            f"{d} = TensorMap('{d}', Interpretation.CATEGORICAL, storage_type=StorageType.CATEGORICAL_FLAG, path_prefix='categorical', "
            f"channel_map={{'no_{d}':0, '{d}':1}}, loss=weighted_crossentropy([1.0, {factor}], '{d}'))\n",
        )
    logging.info(f"Done writing TensorMaps for diseases.")


def _write_disease_tensor_maps_incident_prevalent(phenos_folder: str, f: TextIO) -> None:
    f.write(f"\n\n#  TensorMaps for prevalent and incident MPG disease phenotypes\n")
    disease2tsv = get_disease2tsv(phenos_folder)
    logging.info(f"Got:{len(disease2tsv)} disease TSVs from:{phenos_folder}")
    status_p = disease_prevalence_status(disease2tsv, 1000000, 2500000)
    status_i = disease_incidence_status(disease2tsv, 1000000, 2500000)
    logging.info(f"Got prevalence and incidence status for all diseases.")
    for disease in sorted(list(disease2tsv.keys())):
        total = len(status_p[disease])
        diseased_p = np.sum(list(status_p[disease].values()))
        factor_p = int(total / (1 + (diseased_p * 3)))
        diseased_i = np.sum(list(status_i[disease].values()))
        factor_i = int(total / (1 + (diseased_i * 3)))
        f.write(
            f"{disease}_prevalent_incident = TensorMap('{disease}', Interpretation.CATEGORICAL,  storage_type=StorageType.CATEGORICAL_FLAG, "
            f"path_prefix='categorical', tensor_from_file=prevalent_incident_tensor('dates/enroll_date', 'dates/{disease}_date'), "
            f"channel_map={{'no_{disease}':0, 'prevalent_{disease}':1, 'incident_{disease}':2}}, "
            f"loss=weighted_crossentropy([1.0, {factor_p}, {factor_i}], '{disease}_prevalent_incident'))\n",
        )
    logging.info(f"Done writing TensorMaps for prevalent and incident diseases.")


def _write_phecode_tensor_maps(f: TextIO, phecode_csv, db_client: DatabaseClient):
    # phecode_csv = '/home/sam/phecode_definitions1.2.csv'
    total_samples = 500000
    remove_chars = ";.,/()-[]&' "
    phecode2phenos = {}
    with open(phecode_csv, 'r') as my_csv:
        lol = list(csv.reader(my_csv, delimiter=','))
        for row in lol[1:]:
            pheno = row[1].strip().replace("'s", "s")
            for c in remove_chars:
                pheno = pheno.replace(c, '_')
            pheno = pheno.lower().strip('_').replace('___', '_').replace('__', '_')
            phecode2phenos['phecode_'+row[0].lstrip('0').strip()] = pheno
    query = f"select disease, count(disease) as total from `broad-ml4cvd.ukbb7089_201904.phecodes_nonzero` GROUP BY disease"
    count_result = db_client.execute(query)
    phecode2counts = {}
    for row in count_result:
        phecode2counts[row['disease']] = float(row['total'])

    f.write(f"\n\n#  TensorMaps for Phecode disease phenotypes\n")
    for k, p in sorted(phecode2phenos.items(), key=operator.itemgetter(1)):
        if k in phecode2counts:
            factor = int(total_samples / (1+phecode2counts[k]))
            f.write(
                f"{p}_phe = TensorMap('{k}', Interpretation.CATEGORICAL, channel_map={{'no_{p}':0, '{p}':1}}, path_prefix='categorical', "
                f"storage_type=StorageType.CATEGORICAL_FLAG, loss=weighted_crossentropy([1.0, {factor}], '{k.replace('.', '_')}'))\n",
            )

    query = f"select disease, count(disease) as total from `broad-ml4cvd.ukbb7089_201904.phecodes_nonzero` WHERE prevalent_disease=1 GROUP BY disease"
    count_result = db_client.execute(query)
    phecode2prevalent = {}
    for row in count_result:
        phecode2prevalent[row['disease']] = float(row['total'])

    query = f"select disease, count(disease) as total from `broad-ml4cvd.ukbb7089_201904.phecodes_nonzero` WHERE incident_disease=1 GROUP BY disease"
    count_result = db_client.execute(query)
    phecode2incident = {}
    for row in count_result:
        phecode2incident[row['disease']] = float(row['total'])

    f.write(f"\n\n#  TensorMaps for prevalent and incident Phecode disease phenotypes\n")
    for k, p in sorted(phecode2phenos.items(), key=operator.itemgetter(1)):
        if k in phecode2incident and k in phecode2prevalent:
            factor_i = int(total_samples / (1 + phecode2incident[k]))
            factor_p = int(total_samples / (1 + phecode2prevalent[k]))
            f.write(
                f"{p}_phe_pi = TensorMap('{k}',  Interpretation.CATEGORICAL,  storage_type=StorageType.CATEGORICAL_FLAG, "
                f"path_prefix='categorical', tensor_from_file=prevalent_incident_tensor('dates/enroll_date', 'dates/{k}_date'), "
                f"channel_map={{'no_{p}':0, '{p}_prevalent':1, '{p}_incident':2}}, "
                f"loss=weighted_crossentropy([1.0, {factor_p}, {factor_i}], '{p}_pi'))\n",
            )


def _write_continuous_tensor_maps(f: TextIO, db_client: DatabaseClient):
    # Handle special coding values in continuous variables in order to generate summary statistics (mean and std dev) for
    # each field across all samples. This will remove missing samples from the calculation and change the value of 'Less than one'
    query = f"""
    WITH coding_tmp AS (
        SELECT
            *,
            CASE
                WHEN meaning IN ('Do not know',  'Prefer not to answer', 'Ongoing when data entered') OR meaning LIKE "Still taking%" THEN TRUE
            END AS missing,
            CASE
                WHEN meaning IN {LESS_THAN_CODES} THEN '.5'
            END AS value
        FROM
            {CODING_TABLE}
    ), pheno_tmp AS (
    SELECT
        sample_id,
        FieldID,
        instance,
        array_idx,
        COALESCE(c.value, p.value) new_value,
        COALESCE(c.missing, FALSE) missing
    FROM {PHENOTYPE_TABLE} AS p
    LEFT JOIN coding_tmp AS c
        ON TRUE
        AND SAFE_CAST(p.value AS FLOAT64) = SAFE_CAST(c.coding AS FLOAT64)
        AND p.coding_file_id = c.coding_file_id
    )

    SELECT
        t.FieldID,
        Field,
        t.instance,
        AVG(CAST(new_value AS FLOAT64)) mean,
        STDDEV(CAST(new_value AS FLOAT64)) std,
        MAX(array_idx) AS max_array
    FROM pheno_tmp AS t
    LEFT JOIN {DICTIONARY_TABLE} AS d ON d.FieldID = t.FieldID
    WHERE TRUE
        AND ValueType IN ('Integer', 'Continuous')
        AND NOT missing
    GROUP BY t.FieldID, t.instance, Field ORDER BY t.FieldID
    """

    field_data_for_tensor_maps = db_client.execute(query)

    f.write(f"\n\n#  Continuous tensor maps\n")
    for row in field_data_for_tensor_maps:
        name = dataset_name_from_meaning(None, [str(row.FieldID), row.Field, str(row.instance)])
        channel_map = "channel_map={"
        for i in range(0, row.max_array + 1):
            channel_map += f"'{name}{JOIN_CHAR}{i}': {i}, "
        channel_map += "}"
        f.write(f"ukb_{row.FieldID}_{row.instance} = TensorMap('{name}{JOIN_CHAR}{i}', loss='log_cosh', path_prefix='continuous', ")
        f.write(f"normalization={{'mean': {row.mean}, 'std': {row.std}}}, {channel_map})\n")


def _get_pkl_path_for_field(field_id: int, pyukbb_data_path: str):
    """Returns the path to the .pkl file that contained `UKBioBankParsedField` for a given FieldID."""
    for _, _, files in os.walk(pyukbb_data_path):
        for file in files:
            if file.find(f'FieldID_{field_id}.pkl') > -1:
                return os.path.join(pyukbb_data_path, file)
    raise FileNotFoundError('Cannot find pyukbb .pkl file for field ID {field_id}!')


def _get_all_available_fields(available_fields_pd, keyword: str = None, category: int = None):
    filtered = available_fields_pd
    if category is not None:
        filtered = filtered[filtered.Category == category]
    if keyword is not None:
        filtered = filtered[filtered.Field.str.contains(keyword, case=False)]
    return filtered


def generate_continuous_tensor_map_from_file(
    file_name: str,
    column_name,
    tensor_map_name: str,
    normalization: bool,
    discretization_bounds: List[float],
) -> TensorMap:
    if discretization_bounds:
        return TensorMap(
            f'{tensor_map_name}', Interpretation.DISCRETIZED, channel_map={tensor_map_name: 0},
            tensor_from_file=build_tensor_from_file(file_name, column_name, normalization),
            discretization_bounds=discretization_bounds,
        )
    else:
        return TensorMap(
            f'{tensor_map_name}', channel_map={tensor_map_name: 0},
            annotation_units=64, days_window=0,
            tensor_from_file=build_tensor_from_file(file_name, column_name, normalization),
        )


def generate_categorical_tensor_map_from_file(
    file_name: str,
    column_name: str,
    tensor_map_name: str,
) -> TensorMap:
    ext = file_name.split('.')[1]
    delimiter = ',' if ext == 'csv' else '\t'
    df = pd.read_csv(file_name, delimiter=delimiter)
    channel_map = {}
    for i, k in enumerate(df[column_name].value_counts().keys()):
        channel_map[k] = i
    return TensorMap(
            f'{tensor_map_name}', Interpretation.CATEGORICAL, channel_map=channel_map,
            tensor_from_file=build_categorical_tensor_from_file(file_name, column_name),
    )


def _space_tensor_from_file(df: pd.DataFrame, dimensions: int, sample_column: str = 'sample_id'):
    def tensor_from_file(tm: TensorMap, hd5: h5py.File, dependents=None):
        sample_id = int(os.path.basename(hd5.filename).replace('.hd5', ''))
        row = df[df[sample_column] == sample_id]
        if len(row) == 0:
            raise KeyError(f'Sample id not in dataframe.')
        values = row.iloc[0].tolist()[1:dimensions+1]
        return np.array(values, dtype=np.float32)
    return tensor_from_file


def generate_latent_tensor_map_from_file(
    file_name: str,
    tensor_map_name: str,
) -> TensorMap:
    delimiter = ',' if file_name.split('.')[1].lower() == 'csv' else '\t'
    df = pd.read_csv(file_name, delimiter=delimiter)
    dimensions = len(df.columns)-1
    return TensorMap(
            f'{tensor_map_name}', Interpretation.CONTINUOUS, shape=(dimensions,),
            tensor_from_file=_space_tensor_from_file(df, dimensions), metrics=[],
    )


def generate_random_text_tensor_maps(text_file: str, window_size: int) -> Tuple[TensorMap, TensorMap]:
    name = os.path.basename(text_file).split('.')[0]
    text, token_dictionary = token_dictionary_and_text_from_file(text_file)
    shape = (window_size,)

    output_map = TensorMap(
        f'next_{name}', Interpretation.LANGUAGE,
        shape=shape,
        loss=sparse_cross_entropy(window_size),
        channel_map=token_dictionary,
        cacheable=False,
    )
    input_map = TensorMap(
        name, Interpretation.LANGUAGE, shape=shape,
        tensor_from_file=random_text_window_tensor(text, window_size),
        dependent_map=output_map,
        channel_map=token_dictionary,
        cacheable=False,
    )
    return input_map, output_map


def generate_random_pixel_as_text_tensor_maps(
    tensors: str,
    path_prefix: str,
    window_shape: Tuple[int],
) -> Tuple[TensorMap, TensorMap]:
    name = path_prefix.split('/')[-1]
    path_prefix = '/'.join(path_prefix.split('/')[:-1])
    logging.info(f'prefix {path_prefix} and name is {name}')
    token_dictionary = token_dictionary_from_hd5_key(tensors, path_prefix, name)
    window_size = int(np.prod(window_shape))
    shape = (window_size,)
    output_map = TensorMap(
        f'next_{name}', Interpretation.LANGUAGE,
        shape=shape,
        loss=sparse_cross_entropy(window_size),
        channel_map=token_dictionary,
        cacheable=False,
    )
    input_map = TensorMap(
        name, Interpretation.LANGUAGE, shape=shape, path_prefix=path_prefix,
        tensor_from_file=random_array_window_tensors(window_shape),
        dependent_map=output_map,
        channel_map=token_dictionary,
        cacheable=False,
    )
    return input_map, output_map
