# coding: utf-8
# Copyright (c) 2021 AkaiKKRteam.
# Distributed under the terms of the Apache License, Version 2.0.


from copy import deepcopy
import pandas as pd
import numpy as np
from sklearn.metrics import mean_absolute_error

_SUCCESS_ = "O"
_FAILED_ = "X"
_NO_REF_ = "-"

_FILL_STR_ = "-"
_PAD_STR_ = "  "

_CURRENT_ = "current"
_REFERENCE_ = "reference"

_UPDN_LIST_ = ["up", "dn"]

_UNNUMBERED_TARGET_ = ["te", "tm", "Tc", "resis"]
_NUMBERED_TARGET_ = ["spinlocalmoment", "orbitallocalmoment",
                     "cnd"]


def show_chk_legend():
    print("{}: passed. {}: failed, {}: no reference".format(
        _SUCCESS_, _FAILED_, _NO_REF_))
    print()


def chk_dic_success_all(chk_dic):
    chk_array = np.array([v == _SUCCESS_ for v in chk_dic.values()])
    return chk_array


def _go_fix_localmoment(result, ref, shortname=False):
    """ref can be None"""
    result = deepcopy(result)
    ref = deepcopy(ref)
    df1 = pd.DataFrame([[x for x in result.values()]], index=[_CURRENT_],
                       columns=result.keys())
    if ref is not None:
        df2 = pd.DataFrame([[x for x in result.values()]], index=[_REFERENCE_],
                           columns=result.keys())
    else:
        df2 = None
    return df1, df2
    if True:
        for obj in [result, ref]:
            if obj is None:
                continue
            for label in _NUMBERED_TARGET_:
                if label in obj:
                    v = obj[label]
                    if isinstance(v, list):
                        del obj[label]
                        for i, r in enumerate(v):
                            name = "{}{}".format(label, i+1)
                            obj[name] = r

    result.update({"target": _CURRENT_})
    if ref is not None:
        ref.update({"target": _REFERENCE_})
    df1 = pd.DataFrame(result, dtype=object, index=[0])
    try:
        df2 = pd.DataFrame(ref, dtype=object)
    except ValueError:
        df2 = pd.DataFrame(ref, dtype=object, index=[0])

    # replace long name with short name
    if shortname:
        for df in [df1, df2]:
            _make_df_shortname(df)
    return df1, df2


class ThresBase:
    def __init__(self, thres):
        self.thres = thres
        thres_key_type = []
        thres_key_name = []
        for thres_key, thres_value in thres.items():
            s = thres_key.split("_")
            thres_key_type.append(deepcopy(s[0]))
            thres_key_name.append(deepcopy(s[1]))
        thres_key_type = list(set(thres_key_type))
        thres_key_name = list(set(thres_key_name))
        thres_key_type.sort()
        thres_key_name.sort()
        self.thres_key_type = thres_key_type
        self.thres_key_name = thres_key_name

    @property
    def key_type(self):
        return self.thres_key_type

    @property
    def key_name(self):
        return self.thres_key_name


def _make_df_diff(key, result, ref, thres):
    df1, df2 = _go_fix_localmoment(result, ref)
    if ref is None:
        return df1
    df = pd.concat([df1, df2])
    res = {}
    thresbase = ThresBase(thres)
    for thres_key_type in thresbase.key_type:
        res[thres_key_type] = {}
        thres_key_type_th = thres_key_type+"_th"
        res[thres_key_type_th] = {}

    for thres_key, thres_value in thres.items():
        s = thres_key.split("_")
        thres_key_type = deepcopy(s[0])
        thres_key_name = deepcopy(s[1])
        thres_key_type_th = thres_key_type+"_th"

        for name in df.columns:
            if thres_key_name == name:
                v_current = df.loc[_CURRENT_, name]
                if isinstance(v_current, list):
                    v_current = np.array(v_current)
                v_reference = df.loc[_REFERENCE_, name]
                if isinstance(v_reference, list):
                    v_reference = np.array(v_reference)
                if thres_key_type == "diff":
                    value = np.abs(v_current - v_reference)
                    res[thres_key_type].update({name: value})
                    res[thres_key_type_th].update({name: thres_value})
                elif thres_key_type == "rdiff":
                    value = np.abs(v_current - v_reference) /\
                        np.abs(v_reference)
                    res[thres_key_type].update({name: value})
                    res[thres_key_type_th].update({name: thres_value})

                elif thres_key_type == "and":
                    value = v_current and v_reference
                    res[thres_key_type].update({name: value})
                    res[thres_key_type_th].update({name: thres_value})

                else:
                    print("unkown type", type)
                    print("thres", thres)
                    raise ValueError
    df_list = []
    for res_key, res_value in res.items():
        _df = pd.DataFrame(
            [[v for v in res_value.values()]], columns=res_value.keys(), index=[res_key])
        df_list.append(_df)
    df_comp = pd.concat(df_list, axis=0)
    df_all = pd.concat([df, df_comp], axis=0)
    return df_all


def _make_df_shortname(df):
    return df
    df = df.copy()
    col2 = []
    col = list(df.columns)
    for s in col:
        for lmtarget, abbrev in zip(["spinlocalmoment", "orbitallocalmoment", "threads"],
                                    ["sl", "ol", "thrd"]):
            if lmtarget in s:
                s = s.replace(lmtarget, abbrev)
        col2.append(s)
    df.columns = col2
    return df


def __df_str_add_spc(df, pad=_PAD_STR_):
    lines = []
    for x in df.__str__().splitlines():
        lines.append(pad+x)
    return "\n".join(lines)


class DiffResult:
    def __init__(self, key, result, ref, thres):
        self.key = key
        self.result = result
        self.ref = ref
        self.thres = thres
        df = _make_df_diff(key, result, ref, thres)
        self.df = df

    def process(self, ):
        df = self.df
        thres = self.thres
        if _REFERENCE_ not in df.index:
            return df, None

        res_dic = {}
        for thres_op_label in thres.keys():
            s = thres_op_label.split("_")
            op = s[0]
            label = s[1]
            op_th = op+"_th"
            value = df.loc[op, label]
            if isinstance(value, list):
                value = np.array(value)
            thvalue = df.loc[op_th, label]
            if op == "and":
                flag = value == thvalue
            else:
                flag = value < thvalue
            if isinstance(flag, np.ndarray):
                flag = np.all(flag == True)
            res_dic[thres_op_label] = flag
        df_chk = pd.DataFrame(res_dic, index=["chk"])

        return df, df_chk


class DiffVector:
    def __init__(self, df, thres):
        self.df = df
        self.thres = thres
        self. difflabel = "diff"

    def process(self):
        """make dataframe
        """
        df = self.df
        difflabel = self.difflabel
        v1 = df.loc[_CURRENT_, :].astype(float).values
        v2 = df.loc[_REFERENCE_, :].astype(float).values
        diffvalue = np.abs(v1-v2).tolist()
        _df = pd.DataFrame([v1.tolist(), v2.tolist(), diffvalue], columns=df.columns,
                           index=[_CURRENT_, _REFERENCE_, difflabel])

        self.df = _df  # update self.df
        return _df

    def evaluate_distance(self):
        _df = self.df
        difflabel = self.difflabel
        thres = self.thres

        diff_dic = {}
        for type_ in thres.keys():
            if type_ == "diff_max":
                value = np.max(_df.loc[difflabel, :].values)
            elif type_ == "mae":
                value = mean_absolute_error(_df.loc[_CURRENT_, :].values,
                                            _df.loc[_REFERENCE_, :].values)
            else:
                print("unknown thres.keys()", thres.keys())
                raise ValueError
            diff_dic.update({type_: value})
        df_diff = pd.DataFrame(
            diff_dic, dtype='object', index=["value"])

        th = {}
        for label in ["diff_max", "mae"]:
            th[label] = thres[label]
        df_th = pd.DataFrame(th, index=["thres"])
        df_diff = pd.concat([df_diff, df_th], axis=0)

        chk = {}
        for diffkey in ["diff_max", "mae"]:
            chk[diffkey] = df_diff.loc["value",
                                       diffkey] < df_diff.loc["thres", diffkey]
        df_chk = pd.DataFrame(chk, index=["chk"])

        df_diff = pd.concat([df_diff, df_chk], axis=0)
        return df_diff


def _sort_types_inside_row(df):
    """exchange type1 and type2 as type1<type2 to compare with another data.
    and add pair column.
    """
    df = df.copy()
    type1 = df["type1"].values
    type2 = df["type2"].values
    t1t2_list = []
    for t1, t2 in zip(type1, type2):
        _t1t2 = [t1, t2]
        _t1t2.sort()
        t1t2_list.append(_t1t2)
    df[["type1", "type2"]] = t1t2_list

    if "pair" in list(df.columns):
        del df["pair"]

    # add pair column
    comp1 = df["comp1"].values
    comp2 = df["comp2"].values
    type1 = df["type1"].values
    type2 = df["type2"].values
    typepair = []
    for t1, t2, c1, c2 in zip(type1, type2, comp1, comp2):
        typepair.append("-".join([t1, t2, c1, c2]))
    df_pair = pd.DataFrame({"pair": typepair})

    jijdf = pd.concat([df, df_pair], axis=1)
    return jijdf


def _make_jij_dataframe(result_jij, ref_jij, target="J_ij(meV)"):
    # make jij dataframe
    df_result_jij = pd.DataFrame(result_jij[1:], columns=result_jij[0])
    df_result_jij = df_result_jij[[
        "comp1", "comp2", "J_ij", "J_ij(meV)", "pair"]]

    if ref_jij is not None:
        df_ref_jij = pd.DataFrame(ref_jij[1:], columns=ref_jij[0])
        df_ref_jij = _sort_types_inside_row(df_ref_jij)

        # sort values to compare the result with another dataframe.
        df_ref_jij.sort_values(by="distance", inplace=True)
        df_ref_jij = df_ref_jij[[
            "comp1", "comp2", "J_ij", "J_ij(meV)", "pair"]]
        df_ref_jij.reset_index(drop=True)

        zipped_verson_df_list = zip([_CURRENT_, _REFERENCE_], [
            df_result_jij, df_ref_jij])
    else:
        zipped_verson_df_list = zip([_CURRENT_], [df_result_jij])
        df_ref_jij = None

    # replace J_ij* with {_CURRENT_}_J_ij*.
    # replace distace with {_CURRENT_}_distance.
    for version, _df in zipped_verson_df_list:
        col2 = {}
        for col in _df.columns:
            if col.startswith("J_ij"):
                col2[col] = "{}_{}".format(version, col)
            if col.startswith("distance"):
                col2[col] = "{}_{}".format(version, col)
        _df.rename(columns=col2, inplace=True)

    if True:
        print()
        print("debug df_result_jij")
        print(df_result_jij)

        print("debug df_ref_jij")
        print(df_ref_jij)
        print()

    if ref_jij is not None:
        # merge
        col = ["comp1", "comp2", 'pair']
        df = df_result_jij.merge(df_ref_jij, on=col)
    else:
        df = df_result_jij

    # add pair-comp1-comp2 field
    paircomp = []
    for comp1, comp2, pair in zip(df["comp1"], df["comp2"], df["pair"]):
        paircomp.append("{}_{}_{}".format(pair, comp1, comp2))
    df["typecomp"] = paircomp

    if "{}_{}".format(_REFERENCE_, target) in df.columns:
        col = ["{}_{}".format(_CURRENT_, target), "{}_{}".format(
            _REFERENCE_, target), "typecomp"]
    else:
        col = ["{}_{}".format(_CURRENT_, target), "typecomp"]
    _df = df[col].set_index("typecomp").T

    # rename index
    col2 = {}
    for x in col:
        s = x.split("_")
        col2[x] = s[0]
    _df = _df.rename(index=col2)
    return _df


def _spactra_df(result_totaldos, ref_totaldos,
                updn_list=_UPDN_LIST_, updn_label_list=_UPDN_LIST_):

    if ref_totaldos is not None:
        zipped_version_dos = zip([_CURRENT_, _REFERENCE_],
                                 [result_totaldos, ref_totaldos])
    else:
        zipped_version_dos = zip([_CURRENT_, ],
                                 [result_totaldos, ])
    for version, obj in zipped_version_dos:
        for updn, updnlabel in zip(updn_list, updn_label_list):
            if updnlabel in obj:
                dos = obj[updnlabel]
                if len(dos) > 0:
                    name = "{}_{}".format(version, updn)
                    obj[name] = obj.pop(updnlabel)
                else:
                    del obj[updn]

    df_result_totaldos = pd.DataFrame(result_totaldos)
    if ref_totaldos is not None:
        df_ref_totaldos = pd.DataFrame(ref_totaldos)

        if "energy" in df_ref_totaldos:
            # confirm that their eneriges are the same
            energy_result = df_result_totaldos.loc[:, "energy"].values
            energy_ref = df_ref_totaldos.loc[:, "energy"].values
            if np.all(energy_result == energy_ref):
                df = df_result_totaldos.merge(df_ref_totaldos, on="energy")
            else:
                print(
                    "energies are different between the current calculation and reference.")
                print("no check applied.")
                df = df_result_totaldos
        else:
            print("no energy in reference data. but continue.")
            df = pd.concat([df_result_totaldos, df_ref_totaldos], axis=1)
    else:
        df = df_result_totaldos

    # df is now
    #      energy  _CURRENT__up  _CURRENT__dn  _REFERENCE__up  _REFERENCE__dn
    # 0    -1.495     0.00094     0.00089       0.00094       0.00089
    # 1    -1.485     0.00096     0.00090       0.00096       0.00090
    #
    # split them into up and dn
    df_dos_updn = {}
    for updn in updn_list:
        # search **_up" in columns
        col = []
        for s in df.columns:
            if s.endswith("_"+updn):
                col.append(s)

        if len(col) > 0:
            col.append("energy")
            _df = df.loc[:, col].set_index("energy").T
            index2 = {}
            for s in _df.index:
                s2 = s.replace("_"+updn, "")
                index2[s] = s2
            _df = _df.rename(index=index2)
            df_dos_updn[updn] = _df
    return df_dos_updn


def go_diff_msg(key, result, ref,
                thres={"and_conv": True,
                       "rdiff_te": 1e-11,
                       "diff_tm": 1e-5,
                       "diff_spinlocalmoment": 1e-5,
                       "diff_orbitallocalmoment": 1e-5, }):
    diffresult = DiffResult(key, result, ref, thres)
    df, df_chk = diffresult.process()
    # df.dropna(axis=0, inplace=True)
    df.fillna(_FILL_STR_, inplace=True)

    thresbase = ThresBase(thres)
    col = thresbase.key_name
    print(__df_str_add_spc(df[col]))
    print()
    if _REFERENCE_ not in df.index:
        return _NO_REF_

    print(_PAD_STR_, "summary")
    print(__df_str_add_spc(df_chk.T))
    print()

    chk = df_chk.loc["chk", :].values
    finalchk = np.all(chk == True)
    print(_PAD_STR_, "final:", finalchk)
    if ref is not None:
        if finalchk:
            return _SUCCESS_
        else:
            return _FAILED_
    else:
        return _NO_REF_


def cnd_diff_msg(key, result, ref, thres={"rdiff_resis": 1e-5,
                                          "rdiff_cnd": 1e-5}):
    return go_diff_msg(key, result, ref, thres=thres)


def tc_diff_msg(key, result, ref, thres={"rdiff_Tc": 1e-5}):
    return go_diff_msg(key, result, ref, thres=thres)


def j_diff_msg(key, result, ref,
               thres={"rdiff_Tc": 1e-5, },
               thres_vector={"diff_max": 1e-4, "mae": 1e-4, }):

    target = "jij"
    vector_df_title = "J_ij(meV) at the shortest distance"
    result_jij = result.pop(target, None)
    if ref is not None:
        ref_jij = ref.pop(target, None)
    else:
        ref_jij = None
    # please refer ref_jij, not ref from now on.

    if thres is not None:
        chk1_msg = go_diff_msg(key, result, ref, thres=thres)
        if chk1_msg == _SUCCESS_:
            chk1 = True
        elif chk1_msg == _FAILED_:
            chk1 = False
        elif chk1_msg == _NO_REF_:
            chk1 = True
        else:
            print("unknown return value, chk1", chk1_msg)
            raise ValueError
    else:
        chk1 = True

    df_jij = _make_jij_dataframe(result_jij, ref_jij)

    if _REFERENCE_ in df_jij.index:
        diffvector = DiffVector(df_jij, thres_vector)
        df_jij = diffvector.process()
        print(_PAD_STR_, vector_df_title)
        print(__df_str_add_spc(df_jij))
        print()
    else:
        print(_PAD_STR_, vector_df_title)
        print(__df_str_add_spc(df_jij))
        print()
        return _NO_REF_

    df_jij_chk = diffvector.evaluate_distance()
    print(_PAD_STR_, "{} summary".format(target))
    print(__df_str_add_spc(df_jij_chk))
    print()

    # must add df_tc_chk
    chk2 = np.all(df_jij_chk.loc["chk", :].values == True)
    chk = np.array([chk1, chk2])
    finalchk = np.all(chk == True)
    print(_PAD_STR_, "final:", finalchk)
    if finalchk:
        return _SUCCESS_
    else:
        return _FAILED_


def vector_diff_msg(result_totaldos, ref_totaldos,
                    vector_df_title, thres_vector, updn_label_list=_UPDN_LIST_,):

    df_dos_updn = _spactra_df(
        result_totaldos, ref_totaldos, updn_label_list=updn_label_list)
    if _REFERENCE_ in df_dos_updn["up"].index:
        df_chk_updn = {}
        for (updn, df_dos), vector_df_title_spin in zip(df_dos_updn.items(), vector_df_title):
            diffvector = DiffVector(df_dos, thres_vector)
            df_dos = diffvector.process()
            print(_PAD_STR_, vector_df_title_spin)
            print(__df_str_add_spc(df_dos))
            print()
            if _REFERENCE_ in df_dos.index:
                df_chk = diffvector.evaluate_distance()
                print(__df_str_add_spc(df_chk))
                print()
                df_chk_updn[updn] = df_chk
    else:
        for (updn, df_dos), vector_df_title_spin in zip(df_dos_updn.items(), vector_df_title):
            print(_PAD_STR_, vector_df_title_spin)
            print(__df_str_add_spc(df_dos))
            print()
        return _NO_REF_

    chk_list = []
    for df_chk in df_chk_updn.values():
        chk = df_chk.loc["chk", :].values
        chk_list.append(np.all(chk == True))
    chk_list = np.array(chk_list)
    finalchk = np.all(chk_list == True)
    print(_PAD_STR_, "vector diff:", finalchk)
    print()
    if finalchk:
        return _SUCCESS_
    else:
        return _FAILED_


def dos_diff_msg(key, result, ref, target_list=["totalDOS", "pdos"],
                 updn_label_list=_UPDN_LIST_,
                 thres_vector={"diff_max": 2.5e-3,
                               "mae": 2.5e-4, }):
    result = deepcopy(result)

    chk_list = []
    for target in target_list:
        result_totaldos = result.pop(target)
        if ref is not None:
            ref = deepcopy(ref)
            ref_totaldos = ref.pop(target, None)
        else:
            ref_totaldos = None

        if target == "totalDOS":
            vector_df_title = []
            for updn in updn_label_list:
                vector_df_title.append("{} ({})".format(target, updn))
        elif target == "pdos":
            vector_df_title = []
            atom = result_totaldos["atom"]
            l = result_totaldos["l"]
            for updn in updn_label_list:
                vector_df_title.append(
                    "{} ({}), atom={}, l={}".format(target, updn, atom, l))

        chk = vector_diff_msg(result_totaldos, ref_totaldos,
                              vector_df_title=vector_df_title, thres_vector=thres_vector)
        if chk == _SUCCESS_ or chk == _NO_REF_:
            chk = True
        else:
            chk = False
        chk_list.append(chk)

    chk_list = np.array(chk_list)
    finalchk = np.all(chk_list == True)
    print(_PAD_STR_, "final:", finalchk)
    print()
    if finalchk:
        return _SUCCESS_
    else:
        return _FAILED_


def aw_diff_msg(key, result, ref,
                target="Awk",
                updn_label_list=["Aw_up", "Aw_dn"],
                thres_vector={"diff_max": 3e-3,
                              "mae": 1e-4, }):

    result = deepcopy(result)
    resultAwk = result.pop(target, None)
    if ref is not None:
        ref = deepcopy(ref)
        refAwk = ref.pop(target, None)
    else:
        refAwk = None

    updn_label_list = []
    for updn in _UPDN_LIST_:
        updn_label_list.append("Aw_{}".format(updn))

    vector_df_title = []
    for updn in _UPDN_LIST_:
        vector_df_title.append("Aw(w,k) k={}, {}".format(
            resultAwk["Awk_kpoint"], updn))

    return vector_diff_msg(resultAwk, refAwk, updn_label_list=updn_label_list,
                           vector_df_title=vector_df_title, thres_vector=thres_vector)
