"""Replace any complex model by a linear one, and compute model uncertainties from them."""
import logging
from copy import copy
from functools import reduce
from random import sample

import numpy as np
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import train_test_split
from sympy import symbols, lambdify, diff, simplify
from sympy.matrices import Matrix


# noinspection PyShadowingNames
def linear_error_modeling(data_dict, s1, s2, mdl_vars, logger=None, inp_model=None, n_max_test=500, n_max_fit=1000,
                          min_samples=100):
    """
    :param min_samples:
    :param n_max_fit:
    :param n_max_test:
    :param inp_model:
    :param data_dict: dict with keys: "tau", "sp", "wv", "I", "dtau", "dsp", "dwv", "dI" , "r",
           either ndarrays or scalar (only for differentials)
    :param s1: tuple e.g. ((1,1,1,1),)
    :param s2: tuple e.g. ((0,0,0,0),(0,0,0,0),(0,0,0,0),(0,0,0,0))
    :param mdl_vars: list of variable names used for the modeling, e.g. ["tau", "sp","wv","I"]
    :param logger: either None or logging instance
    :return: dictionary: {"error":[ndarray],"coef":[list],"f":[str, function definition],
             "fit_quality":[tuple, (intercept, offset, norm)]}
    """

    logger = logger if logger is not None else logging.getLogger(__name__)
    dtype = np.float
    # generade differential variables
    mdl_dvars = ["d%s" % var for var in mdl_vars]
    logger.info("Build model based on: %s and derivatives:%s" % (str(mdl_vars), str(mdl_dvars)))
    # dependency correlation and cross correlation
    sel1, sel2 = Matrix(s1), Matrix(s2)
    # define symbols and their derivatives
    vv = symbols(" ".join(mdl_vars))  # list of symbols
    dv = symbols(" ".join(mdl_dvars))  # list of derivative derivatives
    # build model based on correlations and cross correlations
    par = Matrix([vv])
    ff = ((par * sel1.T)[0] + (par * sel2 * par.T)[0]).expand()  # expansion: a_i's
    ww = [symbols("w%i" % ii) for ii in range(len(ff.args) + 1)]  # linear coefficients w_i's
    mdl = ww[0] + sum([arg * w for arg, w in zip(ff.args, ww[1:])])  # final model w0 + sum_i(w_i*a_i)
    mdl_error = simplify(sum([abs(diff(mdl, x) * dx) for x, dx in zip(vv, dv)]))  # gaussian (total) error propagation

    finites = list(
        np.where(  # get list of integer with valid data points
            reduce(np.logical_and,  # all should be true -> reduce list of data with logical_and
                   # generator to produce list of finites in each data field -> to be consumed from logical_and
                   (np.isfinite(v.ravel()) for k, v in data_dict.items() if k in mdl_vars + ["r"])))[0])
    if len(finites) > n_max_fit + n_max_test:
        finites = sample(finites, n_max_fit + n_max_test)
    if len(finites) < min_samples:
        logger.info("Fit failed due to lack of data -> proceed with zero errors.")
        return {
            "error": np.zeros(data_dict["r"].shape, dtype=np.float16),
            "coef": "",
            "f": "",
            "fit_quality": 0.0,
            "y": "",
            "yl": "",
        }
    finites_fit, finites_test = train_test_split(finites,
                                                 test_size=int(n_max_test) / (int(n_max_fit) + int(n_max_test)))
    del finites

    # two options, 1st buld model from inp, 2nd build model from inp_model
    if inp_model is None:
        logger.info("Build model based in inp.")
        inp_deps = {k: np.array(v.ravel()[finites_fit], dtype=dtype) for k, v in data_dict.items() if k in mdl_vars}
        x = np.vstack((lambdify(args=vv, expr=arg, dummify=False)(**inp_deps) for arg in ff.args)).transpose()
        y = np.array(data_dict["r"].ravel()[finites_fit], dtype=dtype)  # reflectance
    else:
        logger.info("Build model based in inp_model.")
        inp_deps = {k: np.array(v.ravel(), dtype=dtype) for k, v in inp_model.items() if k in mdl_vars}
        x = np.vstack([lambdify(args=vv, expr=arg, dummify=False)(**inp_deps) for arg in ff.args]).transpose()
        y = np.array(inp_model["r"].ravel(), dtype=dtype)  # reflectance

    # start fitting
    lr = LinearRegression()
    lr.fit(X=x, y=y)
    # save results
    ww_v = [lr.intercept_] + list(lr.coef_)
    # evaluate model

    # prepare input for sklearn alg,e.g. ff.args=(I, sp, tau, wv, I**2) -> loop over and compute values
    inp_deps = {k: np.array(v.ravel()[finites_test], dtype=dtype) for k, v in data_dict.items() if k in mdl_vars}
    x_test = np.vstack([np.array(
        lambdify(args=vv, expr=arg, dummify=False)(**inp_deps), dtype=dtype) for arg in ff.args]).transpose()
    y_test = np.array(data_dict["r"].ravel()[finites_test], dtype=dtype)  # reflectance
    # noinspection PyArgumentList
    y_model = np.array(lr.predict(X=x_test), dtype=dtype)

    # test model -> scatter plot
    lrt = LinearRegression()
    lrt.fit(X=y_model.reshape(-1, 1), y=y_test.reshape(-1, 1))
    # offset, slope, normalized sum of squares
    fit_quality = (float(lrt.intercept_), float(lrt.coef_[0]), float(np.sum(np.abs(y_model - y_test)) / len(y_test)))
    logger.info("model quality: a0=%.5f,a1=%.5f,<a>=%.5f" % fit_quality)
    logger.info("model: y = %s" % mdl)
    logger.info("error: dy = %s" % mdl_error)
    logger.info("coefficients: w0 -> %.3f" % ww_v[0])
    for w, w_, f in zip(ww[1:], ww_v[1:], ff.args):
        logger.info("coefficients: %s -> %.3f -> %s" % (w, w_, f))
        # prepare to replace coefficients with fitted values
    mdl_error_final = copy(mdl_error)
    mdl_final = copy(mdl)
    for w, v in zip(ww, ww_v):
        mdl_error_final = mdl_error_final.replace(w, v)
        mdl_final = mdl_final.replace(w, v)
    logger.info("error = %s" % mdl_error_final)
    # prepare for numpy evaluation
    f_error = lambdify(args=(vv + dv), expr=mdl_error_final, dummify=False)
    # evaluate model, gather results
    # input values for dependencies and their errors, if numpy array (has ravel() method), then use it,
    # else assume scalar value and proceed
    inp_d_deps = {k: v.ravel() if "ravel" in dir(v) else v for k, v in
                  data_dict.items() if k in mdl_vars + mdl_dvars}

    bf = f_error(**inp_d_deps)
    try:
        error_array = np.array(bf, dtype=np.float16).reshape(data_dict["r"].shape)
    except ValueError:  # in case that bf is only scalar -> get array now
        logger.info("Resulting error field is only scalar.")
        error_array = np.ones(data_dict["r"].shape, dtype=np.float16) * np.float16(bf)
    error_array[np.isnan(data_dict["r"])] = np.NaN

    return {
        "error": error_array,
        "coef": ww_v,
        "f": str(mdl_final),
        "f_error": str(mdl_error_final),
        "fit_quality": fit_quality,
        "y_test": y_test,
        "y_model": y_model
    }


if __name__ == "__main__":
    shape = (5, 5)
    inp = {v: np.random.rand(*shape) for v in ["tau", "sp", "wv", "I", "r", "dtau", "dsp", "dwv", "dI"]}
    inp["dsp"] = 0.5
    inp["dtau"] = 0.1
    mdl_vars = ["tau", "sp", "wv", "I"]
    s1 = ((1, 1, 1, 1),)
    s2 = ((0, 0, 0, 0), (0, 0, 0, 0), (0, 0, 0, 0), (0, 0, 0, 0))

    logger = logging.getLogger("w")
    logger.setLevel(logging.INFO)
    ch = logging.StreamHandler()
    ch.setLevel(logging.INFO)
    logger.addHandler(ch)
    er = linear_error_modeling(data_dict=inp, s1=s1, s2=s2, mdl_vars=mdl_vars, logger=logger)
    print(er)
