import numpy as np
import pandas as pd
import statsmodels.api as sm

from datetime import datetime, timedelta
import yfinance as yf #to download stock price data

import matplotlib.pyplot as plt

from merlion.plot import plot_anoms
from merlion.utils import TimeSeries

np.random.seed(5678)

ticker = '^GSPC'
cur_data = yf.Ticker(ticker)
hist = cur_data.history(period="max")
print(ticker, hist.shape, hist.index.min())

df=hist[hist.index>='2000-01-01'].copy(deep=True)
df.head()

#for each stock_id, get the max close in next 20 trading days
price_col = 'Close'
roll_len=20
new_col = 'next_20day_max'
target_list = []

df.sort_index(ascending=True, inplace=True)
df.head(3)

df_next20dmax=df[[price_col]].shift(1).rolling(roll_len).max()
df_next20dmax.columns=[new_col]
df = df.merge(df_next20dmax, right_index=True, left_index=True, how='inner')

df.dropna(how='any', inplace=True)
df['target']= 100*(df[new_col]-df[price_col])/df[price_col]  

df.head(3)

df.shape

train_ = df[['target']].iloc[:-400].copy(deep=True)
test_ = df[['target']].iloc[-400:].copy(deep=True)

train_data = TimeSeries.from_pd(train_)
test_data = TimeSeries.from_pd(test_)

# Import models & configs
from merlion.models.anomaly.isolation_forest import IsolationForest, IsolationForestConfig
from merlion.models.anomaly.vae import VAE, VAEConfig
from merlion.models.ensemble.anomaly import DetectorEnsemble, DetectorEnsembleConfig
from merlion.post_process.threshold import AggregateAlarms

# isolation forest
iso_forest_config = IsolationForestConfig()
iso_forest_model  = IsolationForest(iso_forest_config)

# VAE
vae_config = VAEConfig()
vae_model  = VAE(vae_config)

#ensemble
en_config = DetectorEnsembleConfig(threshold=AggregateAlarms(alm_threshold=4))
en_model = DetectorEnsemble(config=en_config, models=[iso_forest_model, vae_model])

iso_forest_train_score = iso_forest_model.train(train_data=train_data, anomaly_labels=None)
vae_train_score = vae_model.train(train_data=train_data, anomaly_labels=None)
en_train_score = en_model.train(train_data=train_data, anomaly_labels=None)

df_train_scores = train_.merge(iso_forest_train_score.to_pd(), left_index=True, right_index=True, how='inner')
df_train_scores = df_train_scores.merge(vae_train_score.to_pd(), left_index=True, right_index=True, how='inner')
df_train_scores = df_train_scores.merge(en_train_score.to_pd(), left_index=True, right_index=True, how='inner')
print(df_train_scores.shape, train_.shape)
df_train_scores.head(2)

if_test_scores = iso_forest_model.get_anomaly_score(test_data)
if_test_scores_df = if_test_scores.to_pd()

if_test_labels = iso_forest_model.get_anomaly_label(test_data)
if_test_labels_df = if_test_labels.to_pd()

vae_test_scores = vae_model.get_anomaly_score(test_data)
vae_test_scores_df = vae_test_scores.to_pd()

vae_test_labels = vae_model.get_anomaly_label(test_data)
vae_test_labels_df = vae_test_labels.to_pd()

en_test_scores = en_model.get_anomaly_score(test_data)
en_test_scores_df = en_test_scores.to_pd()

en_test_labels = en_model.get_anomaly_label(test_data)
en_test_labels_df = en_test_labels.to_pd()

df_test_scores = test_.merge(if_test_scores_df, left_index=True, right_index=True, how='inner')
df_test_scores = df_test_scores.merge(vae_test_scores_df, left_index=True, right_index=True, how='inner')
df_test_scores = df_test_scores.merge(en_test_scores_df, left_index=True, right_index=True, how='inner')

df_test_scores = df_test_scores.merge(if_test_labels_df, left_index=True, right_index=True, how='inner')
df_test_scores = df_test_scores.merge(vae_test_labels_df, left_index=True, right_index=True, how='inner')
df_test_scores = df_test_scores.merge(en_test_labels_df, left_index=True, right_index=True, how='inner')
print(test_.shape, df_test_scores.shape)

df_train_scores.columns=['target', 'iso_forest_score', 'vae_score', 'ensemble_score']
df_test_scores.columns=['target', 'iso_forest_score', 'vae_score', 'ensemble_score', 'iso_forest_label', 'vae_label', 'ensemble_label']

df_test_scores.head(3)

df_train_scores.corr()

df_test_scores.corr()

df_test_scores['iso_forest_label'].value_counts()

df_test_scores['vae_label'].value_counts()

import plotly.express as px
import plotly.graph_objects as go
from plotly.subplots import make_subplots

fig_list =[]

# Create figure with secondary y-axis
fig = make_subplots(specs=[[{"secondary_y": True}]])
           


fig.add_trace(go.Scatter(
                        name="target",
                        mode="lines", x=df_train_scores.index,y=df_train_scores['target']),
              secondary_y=False
             )

fig.add_trace(go.Scatter(
                        name="iso_forest_score",
                        mode="lines", x=df_train_scores.index,y=df_train_scores['iso_forest_score']),
              secondary_y=True
             )

fig.add_trace(go.Scatter(
                        name="vae_score",
                        mode="lines", x=df_train_scores.index,y=df_train_scores['vae_score']),
              secondary_y=True
             )


fig.update_layout(hovermode="x unified", 
                  title_text="Merlion Anomaly Detection with Isolation Forest and VAE - training data"
                 )



# Set y-axes titles
fig.update_yaxes(title_text="<b>target</b>", secondary_y=False)
fig.update_yaxes(title_text="<b>anomaly score: </b> isolation forest and vae", secondary_y=True)



fig.update_xaxes(
    title_text="date", 
    rangeslider_visible=True,
    rangeselector=dict(
        buttons=list([
            dict(count=1, label="1m", step="month", stepmode="backward"),
            dict(count=6, label="6m", step="month", stepmode="backward"),
            dict(count=1, label="YTD", step="year", stepmode="todate"),
            dict(count=1, label="1y", step="year", stepmode="backward"),
            dict(step="all")
        ])
    )
)
fig.show()

fig_list.append(fig)

# Create figure with secondary y-axis
fig = make_subplots(specs=[[{"secondary_y": True}]])



fig.add_trace(go.Scatter(
                        name="target",
                        mode="lines", x=df_test_scores.index,y=df_test_scores['target']),
              secondary_y=False
             )

fig.add_trace(go.Scatter(
                        name="iso_forest_score",
                        mode="lines", x=df_test_scores.index,y=df_test_scores['iso_forest_score']),
              secondary_y=True
             )

fig.add_trace(go.Scatter(
                        name="vae_score",
                        mode="lines", x=df_test_scores.index,y=df_test_scores['vae_score']),
              secondary_y=True
             )


fig.update_layout(hovermode="x unified", 
                  title_text="Merlion Anomaly Detection with Isolation Forest and VAE - testing data"
                 )



# Set y-axes titles
fig.update_yaxes(title_text="<b>target</b>", secondary_y=False)
fig.update_yaxes(title_text="<b>anomaly score: </b> isolation forest and vae", secondary_y=True)



fig.update_xaxes(
    title_text="date", 
    rangeslider_visible=True,
    rangeselector=dict(
        buttons=list([
            dict(count=1, label="1m", step="month", stepmode="backward"),
            dict(count=6, label="6m", step="month", stepmode="backward"),
            dict(count=1, label="YTD", step="year", stepmode="todate"),
            dict(count=1, label="1y", step="year", stepmode="backward"),
            dict(step="all")
        ])
    )
)
fig.show()

fig_list.append(fig)

fig_path = r'html/2_Merlion_Isoforest_VAE.html'
fig_list[0].write_html(fig_path)


with open(fig_path, 'a') as f:
    for fig_i in fig_list[1:]:
        f.write(fig_i.to_html(full_html=False, include_plotlyjs='cdn'))