import io
from contextlib import redirect_stderr, redirect_stdout

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from shiny import reactive, render, req, ui

import pyrsm as rsm
import pyrsm.radiant.utils as ru
from pyrsm.utils import intersect


def ui_predict(self, show_ci=True):
    if show_ci:
        ci = (
            ui.input_checkbox(
                "pred_ci",
                "Show conf. intervals",
                value=self.state.get("pred_ci", False),
            ),
            ui.panel_conditional(
                "input.pred_ci == true",
                ui.input_slider(
                    id="conf",
                    label="Confidence level:",
                    min=0,
                    max=1,
                    value=self.state.get("conf", 0.95),
                ),
            ),
        )
    else:
        ci = None

    return ui.panel_conditional(
        "input.tabs == 'Predict'",
        ui.panel_well(
            ui.input_select(
                "pred_type",
                "Prediction input type:",
                ["None", "Data", "Command", "Data & Command"],
                selected=self.state.get("pred_type", "None"),
            ),
            ui.panel_conditional(
                "input.pred_type == 'Data' || input.pred_type == 'Data & Command'",
                ui.input_select(
                    "pred_datasets",
                    "Prediction data:",
                    self.dataset_list,
                    selected=self.state.get("pred_data", None),
                ),
            ),
            ui.panel_conditional(
                "input.pred_type == 'Command' || input.pred_type == 'Data & Command'",
                ru.input_return_text_area(
                    "pred_cmd",
                    "Prediction command:",
                    rows=3,
                    value=self.state.get("pred_cmd", ""),
                    placeholder="Specify a dictionary of values to to use for prediction, e.g., {'carat': 1, 'cut': 'Ideal'}",  # noqa: E501
                ),
            ),
            ci,
        ),
    )


def make_model_inputs(self, input, output, get_data, mod_type="isBin"):
    if isinstance(mod_type, str):

        @output(id="ui_rvar")
        @render.ui
        def ui_rvar():
            isType = get_data()["var_types"][mod_type]
            return ui.input_select(
                id="rvar",
                label="Response Variable:",
                selected=self.state.get("rvar", None),
                choices=isType,
            )

    elif isinstance(mod_type, dict):

        @output(id="ui_rvar")
        @render.ui
        def ui_rvar():
            isType = get_data()["var_types"][mod_type[input.mod_type()]]
            return ui.input_select(
                id="rvar",
                label="Response Variable:",
                selected=self.state.get("rvar", None),
                choices=isType,
            )

    @output(id="ui_evar")
    @render.ui
    def ui_evar():
        vars = get_data()["var_types"]["all"].copy()
        if (input.rvar() is not None) and (input.rvar() in vars):
            del vars[input.rvar()]

        return ui.input_select(
            id="evar",
            label="Explanatory Variables:",
            selected=self.state.get("evar", []),
            choices=vars,
            multiple=True,
            size=min(8, len(vars)),
            selectize=False,
        )


def make_int_inputs(self, input, output, get_data):
    @output(id="ui_interactions")
    @render.ui
    def ui_interactions():
        if len(input.evar()) > 1 and input.show_interactions() is not None:
            choices = []
            nway = int(input.show_interactions())
            isNum = intersect(
                list(input.evar()), list(get_data()["var_types"]["isNum"].keys())
            )
            if len(isNum) > 0:
                # choices += ru.qterms(isNum, nway=int(input.show_interactions()[0]))
                choices += ru.qterms(isNum, nway=nway)
            choices += ru.iterms(input.evar(), nway=nway)
            return ui.input_select(
                id="interactions",
                label=None,
                selected=self.state.get("interactions", None),
                choices=choices,
                multiple=True,
                size=min(8, len(choices)),
                selectize=False,
            )

    @output(id="ui_evar_test")
    @render.ui
    def ui_evar_test():
        choices = {e: e for e in input.evar()}
        if (
            "show_interactions" in input
            and input.show_interactions() is not None
            and int(input.show_interactions()[0]) > 1
        ):
            choices.update({e: e for e in input.interactions()})
        return ui.input_selectize(
            id="evar_test",
            label="Variables to test:",
            selected=self.state.get("evar_test", None),
            choices=choices,
            multiple=True,
        )

    @output(id="ui_incl_evar")
    @render.ui
    def ui_incl_evar():
        if len(input.evar()) > 0:
            return ui.input_select(
                id="incl_evar",
                label="Explanatory variables to include:",
                selected=self.state.get("incl_evar", None),
                choices=input.evar(),
                multiple=True,
                size=min(8, len(input.evar())),
                selectize=False,
            )

    @output(id="ui_incl_interactions")
    @render.ui
    def ui_incl_interactions():
        if len(input.evar()) > 1:
            if "show_interactions" in input and input.show_interactions() is not None:
                nway = int(input.show_interactions())
            else:
                nway = 2
            choices = ru.iterms(input.evar(), nway=nway)
            return ui.input_select(
                id="incl_interactions",
                label="Interactions to include:",
                selected=self.state.get("incl_interactions", None),
                choices=choices,
                multiple=True,
                size=min(8, len(choices)),
                selectize=False,
            )


def make_estimate(
    self,
    input,
    output,
    get_data,
    fun,
    ret,
    module="",
    run=True,
    ec=None,
    debug=True,
):
    if ec is None:

        def estimation_code():
            data_name, code = (get_data()[k] for k in ["data_name", "code"])
            # from https://discord.com/channels/1109483223987277844/1127817202804985917/1129530385429176371
            # does not work with ifelse for some reason
            # inp = ifelse("lev" in input, input.lev(), None)
            if "lev" in input:
                inp = input.lev()
            else:
                inp = None
            if "weights" in input:
                weights = input.weights()
                if weights == "None":
                    weights = None
            else:
                weights = None

            args = {
                "data": f"""{{"{data_name}": {data_name}}}""",
                "rvar": input.rvar(),
                "lev": inp,
                "evar": list(input.evar()),
                "weights": weights,
            }

            if (
                "show_interactions" in input
                and input.show_interactions() is not None
                and int(input.show_interactions()) > 0
                and len(input.interactions()) > 0
            ):
                args["ivar"] = list(input.interactions())

            if ret == "xgb":
                args.update(
                    {
                        "mod_type": input.mod_type(),
                        "n_estimators": input.n_estimators(),
                        "max_depth": input.max_depth(),
                        "min_child_weight": input.min_child_weight(),
                        "learning_rate": input.learning_rate(),
                        "subsample": input.subsample(),
                        "colsample_bytree": input.colsample_bytree(),
                        "random_state": input.random_state(),
                    }
                )
            elif ret == "rf":
                if input.max_features() in ["sqrt", "log2"]:
                    max_features = input.max_features()
                else:
                    max_features = int(input.max_features())

                args.update(
                    {
                        "mod_type": input.mod_type(),
                        "n_estimators": input.n_estimators(),
                        "max_features": max_features,
                        "min_samples_leaf": input.min_samples_leaf(),
                        "max_samples": input.max_samples(),
                        "random_state": input.random_state(),
                    }
                )
            elif ret == "nn":
                args.update(
                    {
                        "mod_type": input.mod_type(),
                        "hidden_layer_sizes": eval(input.hidden_layer_sizes()),
                        "activation": input.activation(),
                        "solver": input.solver(),
                        "alpha": input.alpha(),
                        # "batch_size": input.batch_size(),
                        "learning_rate_init": input.learning_rate_init(),
                        "max_iter": input.max_iter(),
                    }
                )

            args_str = ru.drop_default_args(args, getattr(eval(f"rsm.{module}"), fun))

            if "extra_args" in input and input.extra_args() != "":
                args_str += ", " + input.extra_args()
            return f"""rsm.{module}.{fun}({args_str})""", code

    else:
        estimation_code = ec

    if debug:

        @output(id="show_estimation_code")
        @render.text
        def show_estimation_code():
            out = io.StringIO()
            with redirect_stdout(out), redirect_stderr(out):
                eval(estimation_code()[0])

            return out.getvalue()

    def show_code():
        sc = estimation_code()
        return f"""{sc[1]}\n{ret} = {sc[0]}"""

    if run:

        @reactive.Calc
        @reactive.event(input.run, ignore_none=True)
        def estimate():
            locals()[input.datasets()] = self.datasets[
                input.datasets()
            ]  # get data into local scope

            ec = estimation_code()
            # exec for multi-line
            # https://stackoverflow.com/a/12698067/1974918
            exec(ec[1])
            return eval(ec[0])

    else:

        @reactive.Calc
        def estimate():
            locals()[input.datasets()] = self.datasets[
                input.datasets()
            ]  # get data into local scope

            ec = estimation_code()
            exec(ec[1])
            return eval(ec[0])

    return show_code, estimate


def make_summary(
    self, input, output, session, show_code, estimate, ret, sum_fun, sc=None
):
    if sc is None:

        def summary_code():
            if "controls" in input:
                args = {c: True for c in input.controls()}
            else:
                args = {}
            if (
                "evar_test" in input
                and input.evar_test() is not None
                and len(input.evar_test()) > 0
            ):
                args["test"] = list(input.evar_test())

            args_string = ru.drop_default_args(args, sum_fun)
            return f"""{ret}.summary({args_string})"""

    else:
        summary_code = sc

    @output(id="show_summary_code")
    @render.text
    def show_summary_code():
        cmd = f"""{show_code()}\n{summary_code()}"""
        return ru.code_formatter(cmd, self, input, session, id="copy_summary")

    @output(id="summary")
    @render.text
    def summary():
        out = io.StringIO()
        with redirect_stdout(out), redirect_stderr(out):
            locals()[ret] = estimate()  # get model object into local scope
            eval(summary_code())
        return out.getvalue()


def make_predict(
    self, input, output, session, show_code, estimate, ret, pred_fun, show_ci=False
):
    def predict_code():
        args = {}
        cmd = input.pred_cmd().strip()

        if input.pred_type() == "Data":
            args["data"] = input.pred_datasets()
        elif input.pred_type() == "Command" and not ru.is_empty(cmd):
            args["data"] = None
            args["cmd"] = cmd
        elif input.pred_type() == "Data & Command" and not ru.is_empty(cmd):
            args["data"] = input.pred_datasets()
            args["data_cmd"] = cmd

        if "pred_ci" in input:
            ci = input.pred_ci()
            if ci:
                args.update({"ci": ci, "conf": input.conf()})

        args_string = ru.drop_default_args(
            args, pred_fun, ignore=["data", "cmd", "data_cmd"]
        )
        return f"""{ret}.predict({args_string})"""

    @output(id="show_predict_code")
    @render.text
    def show_predict_code():
        return ru.code_formatter(
            f"""{show_code()}\npred = {predict_code()}""",
            self,
            input,
            session,
            id="copy_predict",
        )

    @output(id="predict")
    @render.data_frame
    def predict():
        if input.pred_type() != "None":
            locals()[ret] = estimate()  # get model object into local scope
            if input.pred_type() in ["Data", "Data & Command"]:
                locals()[input.pred_datasets()] = self.datasets[
                    input.pred_datasets()
                ]  # get prediction data into local scope
            pred = eval(predict_code())

            summary = "Viewing rows {start} through {end} of {total}"
            if pred.shape[0] > 100_000:
                pred = pred[:100_000]
                summary += " (100K rows shown)"

            return render.DataTable(pred.round(3), summary=summary)
        else:
            return None


def make_plot(self, input, output, session, show_code, estimate, ret, pc=None):
    if pc is None:

        def plot_code():
            plots = input.plots()
            if plots == "pred":
                incl = list(input.incl_evar())
                if "incl_interactions" in input:
                    incl_int = list(input.incl_interactions())
                else:
                    incl_int = []
                cmd = f""", incl={incl}"""
                if len(incl_int) > 0:
                    cmd += f""", incl_int={incl_int}"""
                cmd = f"""{ret}.plot(plots="{plots}" {cmd})"""
            elif plots == "corr":
                cmd = f"""{ret}.plot(plots="{plots}", nobs={input.nobs()})"""
            else:
                cmd = f"""{ret}.plot(plots="{plots}")"""
            return cmd

    else:
        plot_code = pc

    @output(id="show_plot_code")
    @render.text
    def show_plot_code():
        plots = input.plots()
        if plots != "None":
            cmd = f"""{show_code()}\n{plot_code()}"""
            return ru.code_formatter(cmd, self, input, session, id="copy_plot")

    @reactive.Calc
    def gen_plot():
        locals()[ret] = estimate()
        eval(f"""{plot_code()}""")
        fig = plt.gcf()
        width, height = fig.get_size_inches()  # Get the size in inches
        return fig, width * 96, height * 96

    @output(id="plot")
    @render.plot
    def plot():
        plots = input.plots()
        if plots != "None":
            return gen_plot()[0]

    @output(id="plot_container")
    @render.ui
    def plot_container():
        req(estimate(), input.plots())
        width, height = gen_plot()[1:]
        aspect_ratio = height / width
        # return ui.output_plot("plot", height=f"{height}px", width=f"{width}px")
        return ui.div(
            {"style": f"width: 100%; padding-bottom: {aspect_ratio * 100}%; position: relative;"},
            ui.div(
                {"style": "position: absolute; top: 0; left: 0; right: 0; bottom: 0;"},
                ui.output_plot("plot", height="100%", width="100%"),
            ),
        )
