# coding=utf-8

import glob
import os
import pickle

import pandas as pd

from pyscenic.utils import modules_from_adjacencies

RESOURCES_FOLDER = "."
RHO_THRESHOLDS = [0.03, 0.1, 0.3]
MODULES_EXT = "modules.dat"
EXP_MTX_EXT = "mtx.tsv"
ADJACENCIES_EXT = "net.csv"


def get_name(fname):
    return os.path.splitext(os.path.basename(fname))[0]


def exp_mtx_fname(name):
    return os.path.join(RESOURCES_FOLDER, "{}.{}".format(name, EXP_MTX_EXT))


def calc_modules(
    adjacencies, exp_mtx, name, rho_dichotomize, rho_threshold=None, mask_dropouts=None
):
    if rho_dichotomize:
        print(
            "{} - {} masking - rho threshold {}".format(
                name, "with" if mask_dropouts else "without", rho_threshold
            )
        )

        out_fname = os.path.join(
            RESOURCES_FOLDER,
            "{}.{}.{}".format(name, rho_threshold, MODULES_EXT)
            if mask_dropouts
            else "{}.{}.no_mask.{}".format(name, rho_threshold, MODULES_EXT),
        )
    else:
        print("{} - all".format(name))

        out_fname = os.path.join(
            RESOURCES_FOLDER, "{}.all.{}".format(name, MODULES_EXT)
        )

    if os.path.isfile(out_fname):
        return

    modules = list(
        modules_from_adjacencies(
            adjacencies,
            exp_mtx,
            rho_dichotomize=rho_dichotomize,
            rho_threshold=rho_threshold,
            rho_mask_dropouts=mask_dropouts,
        )
    )
    print(len(modules))

    with open(out_fname, "wb") as f:
        pickle.dump(modules, f)


def run():
    for fname in glob.glob(
        os.path.join(RESOURCES_FOLDER, "*.{}".format(ADJACENCIES_EXT))
    ):
        name = get_name(fname)
        mtx_fname = exp_mtx_fname(name)

        # Load datasets.
        adjacencies = pd.read_csv(fname)
        exp_mtx = pd.read_csv(mtx_fname, sep="\t", index_col=0).T

        # Calculate modules.
        for rho_threshold in RHO_THRESHOLDS:
            calc_modules(
                adjacencies,
                exp_mtx,
                name,
                rho_dichotomize=True,
                rho_threshold=rho_threshold,
                mask_dropouts=False,
            )
            calc_modules(
                adjacencies,
                exp_mtx,
                name,
                rho_dichotomize=True,
                rho_threshold=rho_threshold,
                mask_dropouts=True,
            )

        calc_modules(adjacencies, exp_mtx, name, rho_dichotomize=False)


if __name__ == "__main__":
    run()
