import os
import re
import json

import pandas as pd
import numpy as np

from celescope.tools.step import Step, s_common
from celescope.vdj.__init__ import CHAINS
import celescope.tools.utils as utils


def report_prepare(df, outdir):

    json_file = outdir + '/.data.json'
    if not os.path.exists(json_file):
        data = {}
    else:
        fh = open(json_file)
        data = json.load(fh)
        fh.close()

    df = df.sort_values('UMI', ascending=False)
    data['CB_num'] = df[df['mark'] == 'CB'].shape[0]
    data['Cells'] = list(df.loc[df['mark'] == 'CB', 'UMI'])
    data['UB_num'] = df[df['mark'] == 'UB'].shape[0]
    data['Background'] = list(df.loc[df['mark'] == 'UB', 'UMI'])

    with open(json_file, 'w') as fh:
        json.dump(data, fh)


@utils.add_log
def count_vdj(args):

    step_name = f"{args.type}_count_vdj"
    step = Step(args, step_name)

    sample = args.sample
    match_dir = args.match_dir
    UMI_min = args.UMI_min
    outdir = args.outdir
    UMI_count_filter1_file = args.UMI_count_filter1_file
    receptor_type = args.type
    iUMI = int(args.iUMI)
    chains = CHAINS[receptor_type]

    # out file
    cell_confident_file = f"{outdir}/{sample}_cell_confident.tsv"
    cell_confident_count_file = f"{outdir}/{sample}_cell_confident_count.tsv"
    clonetypes_file = f"{outdir}/{sample}_clonetypes.tsv"
    match_clonetypes_file = f"{outdir}/{sample}_match_clonetypes.tsv"

    # read file
    df_UMI_count_filter1 = pd.read_csv(UMI_count_filter1_file, sep='\t')
    if (not match_dir) or (match_dir == "None"):
        match_bool = False
    else:
        match_bool = True
    if match_bool:
        match_cell_barcodes, _match_cell_number = utils.read_barcode_file(match_dir)

    cell_summary_row_list = []

    # cell calling：cell calling: keep UMIs >= UMI_min
    df_UMI_sum = df_UMI_count_filter1.groupby(
        ['barcode'], as_index=False).agg({"UMI": "sum"})
    if (UMI_min == "auto"):
        rank = 20
        df_UMI_sum_sorted = df_UMI_sum.sort_values(["UMI"], ascending=False)
        rank_UMI = df_UMI_sum_sorted.iloc[rank, :]["UMI"]
        UMI_min = int(rank_UMI / 10)
    else:
        UMI_min = int(UMI_min)
    df_UMI_cell = df_UMI_sum[df_UMI_sum.UMI >= UMI_min]
    df_UMI_sum["mark"] = df_UMI_sum["UMI"].apply(
        lambda x: "CB" if (x >= UMI_min) else "UB")
    
    # data
    df = df_UMI_sum.sort_values('UMI', ascending=False)
    step.add_data_item(CB_num=df[df['mark'] == 'CB'].shape[0])
    step.add_data_item(Cells=list(df.loc[df['mark'] == 'CB', 'UMI']))
    step.add_data_item(UB_num= df[df['mark'] == 'UB'].shape[0])
    step.add_data_item(Background=list(df.loc[df['mark'] == 'UB', 'UMI']))


    cell_barcodes = set(df_UMI_cell.barcode)
    cell_number = len(cell_barcodes)
    cell_summary_row_list.append({
        "item": "Estimated Number of Cells",
        "count": cell_number,
        "total_count": cell_number,
    })

    # df_UMI_count_filter1 in cell
    df_cell = df_UMI_count_filter1[df_UMI_count_filter1.barcode.isin(
        cell_barcodes)]
    # filter2: cell wtih UMI >= iUMI of identical receptor type and CDR3
    # combinations.
    df_cell_UMI_count_filter2 = df_cell[df_cell.UMI >= iUMI]

    # cell confident
    df_cell_confident = df_cell_UMI_count_filter2[df_cell_UMI_count_filter2["chain"].isin(
        chains)]
    df_cell_confident = df_cell_confident.sort_values("UMI", ascending=False)
    df_cell_confident = df_cell_confident.groupby(
        ["barcode", "chain"], as_index=False).head(1)

    # count
    df_cell_confident_count = df_cell_confident.set_index(["barcode", "chain"])
    df_cell_confident_count = df_cell_confident_count.unstack()
    df_cell_confident_count.columns = [
        '_'.join(col) for col in df_cell_confident_count]
    df_cell_confident_count = df_cell_confident_count.reset_index()
    df_cell_confident_count.fillna(inplace=True, value="NA")

    # clonetypes
    seqs = ["aaSeqCDR3", "nSeqCDR3"]
    cols = []
    for chain in chains:
        for seq in seqs:
            cols.append("_".join([seq, chain]))

    for col in cols:
        if not (col in list(df_cell_confident_count.columns)):
            df_cell_confident_count[col] = "NA"

    df_clonetypes = df_cell_confident_count.copy()

    df_clonetypes = df_clonetypes.groupby(cols, as_index=False).agg({
        "barcode": "count"})
    # put na last
    df_clonetypes.replace('NA', np.nan, inplace=True)
    df_clonetypes.sort_values(["barcode"] + cols, ascending=False, na_position='last', inplace=True)
    df_clonetypes.replace(np.nan, 'NA', inplace=True)

    total_CDR3_barcode_number = sum(df_clonetypes.barcode)
    df_clonetypes["percent"] = df_clonetypes.barcode / \
        total_CDR3_barcode_number * 100
    df_clonetypes["percent"] = df_clonetypes["percent"].apply(
        lambda x: round(x, 2))

    # add clonetype ID
    df_clonetypes = df_clonetypes.reset_index()
    df_clonetypes["clonetype_ID"] = pd.Series(df_clonetypes.index) + 1
    df_clonetypes.drop(columns=["index"], inplace=True)

    # order
    order = ["clonetype_ID"] + cols + ["barcode", "percent"]
    df_clonetypes = df_clonetypes[order]
    df_clonetypes.rename(columns={"barcode": "barcode_count"}, inplace=True)
    # out clonetypes
    df_clonetypes.to_csv(clonetypes_file, sep="\t", index=False)

    if receptor_type == "TCR":

        UMI_col_dic = {"TRA": "UMI_TRA", "TRB": "UMI_TRB"}
        for chain in UMI_col_dic:
            UMI_col_name = UMI_col_dic[chain]
            if UMI_col_name in df_cell_confident_count.columns:
                df_cell_confident_count[UMI_col_name].replace(
                    "NA", 0, inplace=True)
                Median_chain_UMIs_per_Cell = np.median(
                    df_cell_confident_count[UMI_col_name])
            else:
                Median_chain_UMIs_per_Cell = 0
            cell_summary_row_list.append({
                "item": "Median {chain} UMIs per Cell".format(chain=chain),
                "count": Median_chain_UMIs_per_Cell,
                "total_count": np.nan
            })

        df_TRA_TRB = df_cell_confident_count[
            (df_cell_confident_count.aaSeqCDR3_TRA != "NA") &
            (df_cell_confident_count.aaSeqCDR3_TRB != "NA")
        ]
        cell_with_confident_TRA_and_TRB = df_TRA_TRB.shape[0]
        cell_summary_row_list.append({
            "item": "Cell with TRA and TRB",
            "count": cell_with_confident_TRA_and_TRB,
            "total_count": cell_number,
        })

        """
        df cell barcode filter
        intersect cell_barcodes from scRNA-Seq with barcode from TCR seq
        """
        if match_bool:
            cell_with_match_barcode = match_cell_barcodes.intersection(
                cell_barcodes)
            cell_with_match_barcode_number = len(cell_with_match_barcode)

            df_match = df_cell_confident_count[df_cell_confident_count.barcode.isin(
                match_cell_barcodes)]

            df_match_TRA_TRB = df_match[
                (df_match.aaSeqCDR3_TRA != "NA") &
                (df_match.aaSeqCDR3_TRB != "NA")
            ]
            match_cell_with_TRA_and_TRB = df_match_TRA_TRB.shape[0]

            cell_summary_row_list.append({
                "item": "Cell with Barcode Match",
                "count": cell_with_match_barcode_number,
                "total_count": cell_number,
            })
            cell_summary_row_list.append({
                "item": "Cell with Barcode Match, TRA and TRB",
                "count": match_cell_with_TRA_and_TRB,
                "total_count": cell_number,
            })

    # BCR
    elif receptor_type == "BCR":

        UMI_col_dic = {"IGH": "UMI_IGH", "IGL": "UMI_IGL", "IGK": "UMI_IGK"}
        for chain in UMI_col_dic:
            UMI_col_name = UMI_col_dic[chain]
            if UMI_col_name in df_cell_confident_count.columns:
                df_cell_confident_count[UMI_col_name].replace(
                    "NA", 0, inplace=True)
                df_cell_confident_count_over_zero = df_cell_confident_count[
                    df_cell_confident_count[UMI_col_name] > 0
                ]
                Median_chain_UMIs_per_Cell = np.median(
                    df_cell_confident_count_over_zero[UMI_col_name])
            else:
                Median_chain_UMIs_per_Cell = 0
            cell_summary_row_list.append({
                "item": "Median {chain} UMIs per Cell".format(chain=chain),
                "count": Median_chain_UMIs_per_Cell,
                "total_count": np.nan})

        df_heavy_and_light = df_cell_confident_count[
            (df_cell_confident_count.aaSeqCDR3_IGH != "NA") &
            (
                (df_cell_confident_count.aaSeqCDR3_IGL != "NA") |
                (df_cell_confident_count.aaSeqCDR3_IGK != "NA")
            )
        ]
        Cell_with_Heavy_and_Light_Chain = df_heavy_and_light.shape[0]
        cell_summary_row_list.append({
            "item": "Cell with Heavy and Light Chain",
            "count": Cell_with_Heavy_and_Light_Chain,
            "total_count": cell_number
        })

        """
        df cell barcode filter
        intersect cell_barcodes from normal scRNA-Seq with barcode from BCR seq
        """
        if match_bool:
            cell_with_match_barcode = match_cell_barcodes.intersection(
                cell_barcodes)
            cell_with_match_barcode_number = len(cell_with_match_barcode)

            df_match = df_cell_confident_count[df_cell_confident_count.barcode.isin(
                match_cell_barcodes)]

            # median match UMI
            df_match_heavy_light = df_match[
                (df_match.aaSeqCDR3_IGH != "NA") &
                (
                    (df_match.aaSeqCDR3_IGL != "NA") |
                    (df_match.aaSeqCDR3_IGK != "NA")
                )
            ]
            match_cell_with_heavy_and_light = df_match_heavy_light.shape[0]

            cell_summary_row_list.append({
                "item": "Cell with Barcode Match ",
                "count": cell_with_match_barcode_number,
                "total_count": cell_number
            })
            cell_summary_row_list.append({
                "item": "Cell with Barcode Match, Heavy and Light Chain",
                "count": match_cell_with_heavy_and_light,
                "total_count": cell_number
            })

    if match_bool:
        """
        df_match_clonetypes
        """
        df_match_clonetypes = df_match.groupby(cols, as_index=False).agg({
            "barcode": "count"})
        total_match_CDR3_barcode_number = sum(
            df_match_clonetypes.barcode)
        df_match_clonetypes["percent"] = df_match_clonetypes.barcode / \
            total_match_CDR3_barcode_number * 100
        df_match_clonetypes["percent"] = df_match_clonetypes["percent"].apply(
            lambda x: round(x, 2)
        )
        df_match_clonetypes.rename(columns={"barcode": "barcode_count"}, inplace=True)
        df_match_clonetypes = df_match_clonetypes.merge(
            df_clonetypes, on=cols, how='left', suffixes=('', '_y'))
        # order and drop duplicated cols
        order = ["clonetype_ID"] + cols + ["barcode_count", "percent"]
        df_match_clonetypes = df_match_clonetypes[order]
        df_match_clonetypes.sort_values(["barcode_count", "clonetype_ID"], ascending=[False,True], inplace=True)
        df_match_clonetypes.to_csv(
            match_clonetypes_file, sep="\t", index=False)

    df_mergeID = pd.merge(df_cell_confident_count,
                          df_clonetypes, how="left", on=cols)
    df_mergeID.sort_values(["clonetype_ID", "barcode"], inplace=True)
    # output df_cell_confident_count
    df_mergeID.to_csv(cell_confident_count_file, sep="\t", index=False)
    df_mergeID = df_mergeID[["barcode", "clonetype_ID"]]
    df_cell_confident_with_ID = pd.merge(
        df_cell_confident, df_mergeID, how="left", on="barcode")
    df_cell_confident_with_ID.sort_values(
        ["clonetype_ID", "barcode", "chain"], inplace=True)
    # output df_cell_confident
    df_cell_confident_with_ID.to_csv(
        cell_confident_file, sep="\t", index=False)

    # summary file
    cell_summary = pd.DataFrame(cell_summary_row_list, columns=[
                                "item", "count", "total_count"])
    cell_summary["count"] = cell_summary["count"].apply(int)
    cell_summary["percent"] = cell_summary["count"] / \
        (cell_summary.total_count.astype("float")) * 100
    cell_summary["percent"] = cell_summary["percent"].apply(
        lambda x: round(x, 2))
    cell_summary["count"] = cell_summary["count"].apply(utils.format_number)

    def percent_str_func(row):
        need_percent = bool(
            re.search("Cell with", row["item"], flags=re.IGNORECASE))
        if need_percent:
            return "(" + str(row["percent"]) + "%)"
        else:
            return ""
    cell_summary["percent_str"] = cell_summary.apply(percent_str_func, axis=1)

    # stat file
    def gen_stat(summary, stat_file):
        stat = summary
        stat["new_count"] = stat["count"].astype(str) + stat["percent_str"]
        stat = stat.loc[:, ["item", "new_count"]]
        stat.to_csv(stat_file, sep=":", header=None, index=False)

    cell_stat_file = "{}/stat.txt".format(outdir)
    gen_stat(cell_summary, cell_stat_file)
    step.add_data_item(iUMI=iUMI)

    # cloneytpes table
    def format_table(df_clonetypes):
        df_table = df_clonetypes.copy()
        df_table["percent"] = df_table["percent"].apply(
            lambda x: str(x) + "%")
        seqs = ["aaSeqCDR3"]
        cols = []
        for chain in chains:
            for seq in seqs:
                cols.append("_".join([seq, chain]))
        df_table_cols = ["clonetype_ID"] + cols + ["barcode_count", "percent"]
        df_table = df_table[df_table_cols]
        table_header = ["Clonetype_ID"] + cols + ["Frequency", "Percent"]
        return df_table, table_header

    df_table, _table_header = format_table(df_clonetypes)
    title = 'Clonetypes'
    if match_bool:
        df_table, _table_header = format_table(df_match_clonetypes)
        title = 'Match Clonetypes'

    table_dict = step.get_table(title, 'clonetypes_table', df_table)
    step.add_data_item(table_dict=table_dict)
    
    # other_metrics_file
    """
    if len(other_metrics_row_list) != 0:
        other_metrics = pd.DataFrame(other_metrics_row_list,columns=["item","count"])
        other_metrics.to_csv(other_metrics_file,sep=":",header=None,index=False)
    """

    step.clean_up()


def get_opts_count_vdj(parser, sub_program):
    if sub_program:
        parser = s_common(parser)
        parser.add_argument("--UMI_count_filter1_file", required=True)
        parser.add_argument("--match_dir", default=None)
    parser.add_argument("--type", required=True)
    parser.add_argument('--UMI_min', dest='UMI_min',
                        help='minimum UMI number to filter', default="auto")
    parser.add_argument(
        '--iUMI', help='minimum number of UMI of identical receptor type and CDR3', default=1)
