from datetime import timedelta, datetime, date
from unittest.mock import patch
from typing import Callable, Any
from inspect import signature
import streamlit as st
import pandas as pd
import yfinance as yf
from openbb_terminal.sdk import openbb
from openbb_terminal.forecast import helpers
from openbb_terminal.rich_config import console

st.set_page_config(layout="wide")

# pylint: disable=E1101
model_opts = {
    "expo": openbb.forecast.expo,  # type: ignore
    "theta": openbb.forecast.theta,  # type: ignore
    "linregr": openbb.forecast.linregr,  # type: ignore
    "regr": openbb.forecast.regr,  # type: ignore
    "rnn": openbb.forecast.rnn,  # type: ignore
    "brnn": openbb.forecast.brnn,  # type: ignore
    "nbeats": openbb.forecast.nbeats,  # type: ignore
    "tcn": openbb.forecast.tcn,  # type: ignore
    "trans": openbb.forecast.trans,  # type: ignore
    "tft": openbb.forecast.tft,  # type: ignore
}

feat_engs = {
    "ema": openbb.forecast.ema,  # type: ignore
    "sto": openbb.forecast.sto,  # type: ignore
    "rsi": openbb.forecast.rsi,  # type: ignore
    "roc": openbb.forecast.roc,  # type: ignore
    "mom": openbb.forecast.mom,  # type: ignore
    "atr": openbb.forecast.atr,  # type: ignore
    "delta": openbb.forecast.delta,  # type: ignore
    "signal": openbb.forecast.signal,  # type: ignore
}
# pylint: enable=E1101

# Add these: "1m", "2m", "5m", "15m", "30m", "60m", "90m", "1h",

interval_opts = ["1d", "5d", "1wk", "1mo", "3mo"]


def format_df(df: pd.DataFrame) -> pd.DataFrame:
    if len(df.columns) != 6:
        df.columns = ["_".join(col).strip() for col in df.columns.values]
    df.reset_index(inplace=True)
    df.columns = [x.lower() for x in df.columns]
    return df


def has_parameter(func: Callable[..., Any], parameter: str) -> bool:
    params = signature(func).parameters
    parameters = params.keys()
    return parameter in parameters


def load_state(name: str, default: Any):
    if name not in st.session_state:
        st.session_state[name] = default


def special_st(text: str):
    if "[green]" not in text:
        st.write(text)


def run_forecast(
    data: pd.DataFrame,
    model: str,
    target_column: str,
    past_covariates: list[str],
    n_predict: int,
):
    if helpers.check_data(data, target_column):
        # TODO: let the user choose their own n_predict
        kwargs: dict[str, Any] = {}
        forecast_model = model_opts[model]
        contains_covariates = has_parameter(forecast_model, "past_covariates")
        if contains_covariates and past_covariates != []:
            kwargs["past_covariates"] = ",".join(past_covariates)
        if has_parameter(forecast_model, "output_chunk_length"):
            kwargs["output_chunk_length"] = n_predict

        # n_predict and output_chunk_length must be the same if there are past covariates
        # run a spinner while we wait for the model to run
        with st.spinner("Running model..."):
            response = forecast_model(
                data=data,
                target_column=target_column,
                n_predict=n_predict,
                **kwargs,
            )
        if model == "theta":
            (
                ticker_series,
                historical_fcast,
                predicted_values,
                precision,
                _,
                _model,
            ) = response
        else:
            (
                ticker_series,
                historical_fcast,
                predicted_values,
                precision,
                _model,
            ) = response
        del precision
        if model in ["expo", "linregr", "rnn", "tft"]:
            predicted_values = predicted_values.quantile_df()[f"{target_column}_0.5"]
        else:
            predicted_values = predicted_values.pd_dataframe()[target_column]
        return (
            historical_fcast.pd_dataframe(),
            ticker_series.pd_dataframe(),
            pd.DataFrame(predicted_values),
        )
    return pd.DataFrame(), pd.DataFrame(), pd.DataFrame()


class Handler:
    def __init__(self):
        load_state("last_tickers", "")
        load_state("last_intervals", "1d")
        load_state("df", pd.DataFrame())
        default_opts = {
            key: [] for key in ["target_widget", "column_widget", "past_covs_widget"]
        }
        load_state("widget_options", default_opts)

        self.feature_model = None
        self.feature_target = None

    # pylint: disable=R0913
    def handle_changes(
        self,
        past_covariates: list[str],
        start,
        end,
        interval,
        tickers,
        target_column,
        model,
        n_predict,
        naive,
        forecast_only,
    ):
        del naive, forecast_only
        if tickers and target_column:
            start_n = datetime(start.year, start.month, start.day)
            end_n = datetime(end.year, end.month, end.day)
            if interval in ["1d", "5d", "1wk", "1mo", "3mo"]:
                result = st.session_state["df"].loc[
                    (st.session_state["df"]["date"] >= start_n)
                    & (st.session_state["df"]["date"] <= end_n)
                ]
            else:
                result = st.session_state["df"]
            if not target_column:
                target_column = st.session_state["df"].columns[0]
            with patch.object(console, "print", special_st):
                if helpers.check_data(result, target_column):
                    final_df = helpers.clean_data(result)
                    hist_fcast, tick_series, pred_vals = run_forecast(
                        final_df, model, target_column, past_covariates, n_predict
                    )
                    hist_fcast.columns = ["Historical Forecast"]
                    tick_series.columns = ["Past Prices"]
                    pred_vals.columns = ["Prediction Values"]
                    final = pd.concat(
                        [hist_fcast, tick_series, pred_vals], axis=1, join="outer"
                    )
                    if not final.empty:
                        if helpers.check_dates(pred_vals.index.to_series()):
                            pred_vals.index = pred_vals.index.date
                        rowc1, rowc2 = st.columns([4, 1])
                        # Styled write() with sig figs
                        with rowc2:
                            st.write(
                                pred_vals.style.format({"Prediction Values": "{:.2f}"})
                            )
                        with rowc1:
                            st.line_chart(final)

                    else:
                        st.write("There was an error with the data")

    def handle_eng(self, target, feature):
        self.feature_target = target
        self.feature_model = feat_engs[feature]

    def on_ticker_change(self):
        tickers = st.session_state.ticker
        if tickers:
            interval = st.session_state.interval
            start = st.session_state.start
            end = st.session_state.end
            if (
                tickers != st.session_state["last_tickers"]
                or interval != st.session_state["last_interval"]
            ):
                if interval in ["1d", "5d", "1wk", "1mo", "3mo"]:
                    df = yf.download(
                        tickers, period="max", interval=interval, progress=False
                    )
                else:
                    end_date = end + timedelta(days=1)
                    df = yf.download(
                        tickers,
                        start=start,
                        end=end_date,
                        interval=interval,
                        progress=False,
                    )
                df = df.dropna()
                st.session_state["df"] = format_df(df)
                st.session_state["last_tickers"] = tickers
                st.session_state["last_interval"] = interval
        st.session_state["widget_options"]["target_widget"] = [
            x for x in st.session_state["df"] if x != "date"
        ]
        st.session_state["widget_options"]["past_covs_widget"] = [
            x for x in st.session_state["df"] if x != "date"
        ]
        st.session_state["widget_options"]["column_widget"] = [
            x for x in st.session_state["df"] if x != "date"
        ]

    def run(self):

        st.title("OpenBB Forecasting")  # Title does not like being in a column

        r1c1, r1c2, r1c3, r1c4, r1c5 = st.columns([2, 1, 1, 1, 1])
        r2c1, r2c2, r2c3 = st.columns([1, 1, 1])

        with r1c1:
            ticker = st.text_input(
                "Ticker", "", key="ticker", on_change=self.on_ticker_change
            )
        with r1c2:
            start_date = st.date_input(
                "Start date", date.today() - timedelta(weeks=104), key="start"
            )
        with r1c3:
            end_date = st.date_input("End date", date.today(), key="end")
        with r1c4:
            target_interval = st.selectbox(
                "Interval", index=0, key="interval", options=interval_opts
            )
        with r1c5:
            n_predict = st.selectbox(
                "Prediction Days", index=3, key="n_predict", options=list(range(2, 31))
            )
        with r2c1:
            # TODO: disable this if the current model does not allow for it
            past_covs_widget = st.multiselect(
                "Past Covariates",
                options=st.session_state["widget_options"]["past_covs_widget"],
            )
        with r2c2:
            target_widget = st.selectbox(
                "Target", options=st.session_state["widget_options"]["target_widget"]
            )
        with r2c3:
            model_widget = st.selectbox("Model", options=list(model_opts))

        st.markdown("""---""")
        if st.button("Get forecast"):
            if ticker:
                self.handle_changes(
                    past_covariates=past_covs_widget,
                    start=start_date,
                    end=end_date,
                    interval=target_interval,
                    tickers=ticker,
                    target_column=target_widget,
                    model=model_widget,
                    n_predict=n_predict,
                    naive=False,
                    forecast_only=False,
                )
            else:
                st.write("Please select a ticker")

            # TODO need to new add in business days and merge to make new DF before plotting
            # create df with 5 next business days from a start date
            # next_business_days = pd.date_range(start=end_date, periods=5, freq="B")

        # if st.button("Add Column"):
        # kwargs = {}
        # if has_parameter(handler.feature_model, "target_column"):
        # kwargs["target_column"] = handler.feature_target
        # handler.df = handler.feature_model(handler.df, **kwargs)
        # st.session_state["widget_options"]["past_covs_widget"] = handler.df.columns


if __name__ == "__main__":
    handler = Handler()
    handler.run()
