#!/usr/bin/env python3

import csv
import warnings

import mngs
import numpy as np
import pandas as pd
import scipy

if "general" in __file__:
    with warnings.catch_warnings():
        warnings.simplefilter("always")
        warnings.warn(
            '\n"mngs.general.save" will be removed. '
            'Please use "mngs.io.save" instead.',
            PendingDeprecationWarning,
        )


def save(obj, sfname_or_spath, makedirs=True, show=True, **kwargs):
    """
    Example
      save(arr, 'data.npy')
      save(df, 'df.csv')
      save(serializable, 'serializable.pkl')
    """
    import inspect
    import json
    import os
    import pickle

    import h5py
    import numpy as np
    import pandas as pd
    import torch
    import yaml

    spath, sfname = None, None

    if "/" in sfname_or_spath:
        spath = sfname_or_spath
    else:
        sfname = sfname_or_spath

    if (spath is None) and (sfname is not None):
        ## for ipython
        __file__ = inspect.stack()[1].filename
        if "ipython" in __file__:
            __file__ = f'/tmp/fake-{os.getenv("USER")}.py'

        ## spath
        fpath = __file__
        fdir, fname, _ = mngs.general.split_fpath(fpath)
        sdir = fdir + fname + "/"
        spath = sdir + sfname
        # spath = mk_spath(sfname, makedirs=True)

    ## Make directory
    if makedirs:
        sdir = os.path.dirname(spath)
        os.makedirs(sdir, exist_ok=True)

    ## Saves
    try:
        ## copy files
        is_copying_files = (
            isinstance(obj, str) or is_listed_X(obj, str)
        ) and (isinstance(spath, str) or is_listed_X(spath, str))
        if is_copying_files:
            mngs.general.copy_files(obj, spath)

        # csv
        elif spath.endswith(".csv"):
            if isinstance(obj, pd.DataFrame):  # DataFrame
                obj.to_csv(spath)
            if is_listed_X(obj, [int, float]):  # listed scalars
                save_listed_scalars_as_csv(
                    obj,
                    spath,
                    **kwargs,
                )
            if is_listed_X(obj, pd.DataFrame):  # listed DataFrame
                # save_listed_dfs_as_csv(obj, spath, indi_suffix=None, overwrite=False)
                save_listed_dfs_as_csv(obj, spath, **kwargs)
        # numpy
        elif spath.endswith(".npy"):
            np.save(spath, obj)
        # pkl
        elif spath.endswith(".pkl"):
            with open(spath, "wb") as s:  # 'w'
                pickle.dump(obj, s)
        # joblib
        elif spath.endswith(".joblib"):
            with open(spath, "wb") as s:  # 'w'
                joblib.dump(obj, s, compress=3)
        # png
        elif spath.endswith(".png"):
            # plotly
            if str(type(obj)) == "<class 'plotly.graph_objs._figure.Figure'>":
                obj.write_image(file=spath, format="png")
            # matplotlib
            else:
                try:
                    obj.savefig(spath)
                except:
                    obj.figure.savefig(spath)
            del obj
        # tiff
        elif spath.endswith(".tiff") or spath.endswith(".tif"):
            obj.savefig(
                spath, dpi=300, format="tiff"
            )  # obj is matplotlib.pyplot object
            del obj
        # mp4
        elif spath.endswith(".mp4"):
            mk_mp4(obj, spath)  # obj is matplotlib.pyplot.figure object
            del obj

        # yaml
        elif spath.endswith(".yaml"):
            from ruamel.yaml import YAML

            yaml = YAML()
            yaml.preserve_quotes = (
                True  # Optional: if you want to preserve quotes
            )
            yaml.indent(
                mapping=4, sequence=4, offset=4
            )  # Optional: set indentation

            with open(spath, "w") as f:
                yaml.dump(obj, f)

        # json
        elif spath.endswith(".json"):
            with open(spath, "w") as f:
                json.dump(obj, f, indent=4)

        # hdf5
        elif spath.endswith(".hdf5"):
            name_list, obj_list = []
            for k, v in obj.items():
                name_list.append(k)
                obj_list.append(v)
            with h5py.File(spath, "w") as hf:
                for (name, obj) in zip(name_list, obj_list):
                    hf.create_dataset(name, data=obj)
        # pth
        elif spath.endswith(".pth"):
            torch.save(obj, spath)

        # mat
        elif spath.endswith(".mat"):
            scipy.io.savemat(spath, obj)

        # catboost model
        elif spath.endswith(".cbm"):
            obj.save_model(spath)

        else:
            raise ValueError("obj was not saved.")

    except Exception as e:
        print(e)

    else:
        if show and not is_copying_files:
            # if show:
            print(f"\nSaved to: {spath}\n")


def check_encoding(file_path):
    from chardet.universaldetector import UniversalDetector

    detector = UniversalDetector()
    with open(file_path, mode="rb") as f:
        for binary in f:
            detector.feed(binary)
            if detector.done:
                break
    detector.close()
    enc = detector.result["encoding"]
    return enc


def is_listed_X(obj, types):
    """
    Example:
        obj = [3, 2, 1, 5]
        is_listed_X(obj,
    """
    import numpy as np

    try:
        conditions = []
        condition_list = isinstance(obj, list)

        if not (isinstance(types, list) or isinstance(types, tuple)):
            types = [types]

        _conditions_susp = []
        for typ in types:
            _conditions_susp.append(
                (np.array([isinstance(o, typ) for o in obj]) == True).all()
            )

        condition_susp = np.any(_conditions_susp)

        is_listed_X = np.all([condition_list, condition_susp])
        return is_listed_X

    except:
        return False


def save_listed_scalars_as_csv(
    listed_scalars,
    spath_csv,
    column_name="_",
    indi_suffix=None,
    round=3,
    overwrite=False,
    show=False,
):
    """Puts to df and save it as csv"""
    import numpy as np

    if overwrite == True:
        mv_to_tmp(spath_csv, L=2)
    indi_suffix = (
        np.arange(len(listed_scalars)) if indi_suffix is None else indi_suffix
    )
    df = pd.DataFrame(
        {"{}".format(column_name): listed_scalars}, index=indi_suffix
    ).round(round)
    df.to_csv(spath_csv)
    if show:
        print("\nSaved to: {}\n".format(spath_csv))


def save_listed_dfs_as_csv(
    listed_dfs,
    spath_csv,
    indi_suffix=None,
    overwrite=False,
    show=False,
):
    """listed_dfs:
        [df1, df2, df3, ..., dfN]. They will be written vertically in the order.

    spath_csv:
        /hoge/fuga/foo.csv

    indi_suffix:
        At the left top cell on the output csv file, '{}'.format(indi_suffix[i])
        will be added, where i is the index of the df.On the other hand,
        when indi_suffix=None is passed, only '{}'.format(i) will be added.
    """
    import numpy as np

    if overwrite == True:
        mv_to_tmp(spath_csv, L=2)

    indi_suffix = (
        np.arange(len(listed_dfs)) if indi_suffix is None else indi_suffix
    )
    for i, df in enumerate(listed_dfs):
        with open(spath_csv, mode="a") as f:
            f_writer = csv.writer(f)
            i_suffix = indi_suffix[i]
            f_writer.writerow(["{}".format(indi_suffix[i])])
        df.to_csv(spath_csv, mode="a", index=True, header=True)
        with open(spath_csv, mode="a") as f:
            f_writer = csv.writer(f)
            f_writer.writerow([""])
    if show:
        print("Saved to: {}".format(spath_csv))


def mk_mp4(fig, spath_mp4):
    from matplotlib import animation

    axes = fig.get_axes()

    def init():
        return (fig,)

    def animate(i):
        for ax in axes:
            ax.view_init(elev=10.0, azim=i)
        return (fig,)

    anim = animation.FuncAnimation(
        fig, animate, init_func=init, frames=360, interval=20, blit=True
    )

    writermp4 = animation.FFMpegWriter(
        fps=60, extra_args=["-vcodec", "libx264"]
    )
    anim.save(spath_mp4, writer=writermp4)
    print("\nSaving to: {}\n".format(spath_mp4))


def mv_to_tmp(fpath, L=2):
    import os
    from shutil import move

    try:
        tgt_fname = connect_strs_with_hyphens(fpath.split("/")[-L:])
        tgt_fpath = "/tmp/{}".format(tgt_fname)
        move(fpath, tgt_fpath)
        print("Moved to: {}".format(tgt_fpath))
    except:
        pass


def save_optuna_study_as_csv_and_pngs(study, sdir):
    import optuna

    ## Trials DataFrame
    trials_df = study.trials_dataframe()

    ## Figures
    hparams_keys = list(study.best_params.keys())
    slice_plot = optuna.visualization.plot_slice(study, params=hparams_keys)
    contour_plot = optuna.visualization.plot_contour(
        study, params=hparams_keys
    )
    optim_hist_plot = optuna.visualization.plot_optimization_history(study)
    parallel_coord_plot = optuna.visualization.plot_parallel_coordinate(
        study, params=hparams_keys
    )
    hparam_importances_plot = optuna.visualization.plot_param_importances(
        study
    )
    figs_dict = dict(
        slice_plot=slice_plot,
        contour_plot=contour_plot,
        optim_hist_plot=optim_hist_plot,
        parallel_coord_plot=parallel_coord_plot,
        hparam_importances_plot=hparam_importances_plot,
    )

    ## Saves
    save(trials_df, sdir + "trials_df.csv")

    for figname, fig in figs_dict.items():
        save(fig, sdir + f"{figname}.png")
