import pandas as pd
pd.core.common.is_list_like = pd.api.types.is_list_like
import numpy as np
from dfply import *
from plotnine import *
import ipdb
import logging
import pyfolio as pf
from dataclasses import dataclass
from collections import defaultdict
import itertools
import pickle 
from joblib import Parallel
from joblib import delayed
import itertools as it

pd.set_option('display.max_columns', None)


def main():
    ps, ps_summary = setup_and_run()

def get_model_data():
    #model_data = pd.read_feather( '../../data/backtesting/pop_5_base_model.feather')
    model_data = pd.read_feather( '../../data/backtesting/first_seen_dummy.feather')
    model_data = (model_data >> 
            rename( spyid = 'index', streams = 'actual_streams') >> 
            mutate( streams = exp(X.streams), predicted_streams = exp(X.predicted_streams)))

    return model_data

@dataclass
class PortfolioSummary:
    portfolio_stats: pd.DataFrame
    portfolio_summary: pd.DataFrame


@dataclass
class BacktestRun:
    parameters: dict
    ps: PortfolioSummary

    def to_summary_df():
        return pd.concat(parameters, ps.portfolio_stats, axis=1)

def setup_and_run_many():

    # logging.basicConfig(format='[ %(levelname)s ] %(asctime)s -  %(message)s', level=logging.DEBUG)
    logging.basicConfig(format='%(asctime)s - %(message)s', level=logging.WARNING)
    simulations_per_setup = 100
    br_runs = []
    
    setup = { 'payment_per_stream' : [0.005],
            'artist_advance_fraction' : [0.8],
            'pop_5_slippage' : [0.15], # Applied on both enter and exit
            'execution_slippage' : [0.1], # Applied only on the exit, ie how long does it take to sign up the artist
            'clip_individual_returns' : [3],
            'min_price' : [2e3,5e3],
            'max_price' : [7e3,10e3],
            'duration_days' :[365],
            'new_streams_per_date' : [1/7, 3/7, 1, 2],
            'capital' : [250e3] }

    # with Parallel(n_jobs = 3, verbose = 10, prefer='threads') as pl:
        # for cc in capital_amounts:
            # rr = pl(delayed(runOneBacktest)(ss, cc, setup) for ss in range(simulations_per_setup))
            # br_runs.append(rr)


    all_settings = sorted(setup)
    exploded_setups = list(map(lambda x: dict(zip(all_settings, x)), it.product(*(setup[name] for name in all_settings))))

    settings_that_vary = {k for k, v in setup.items() if len(v) > 1}

    import time
    stime = time.time()

    sim = 0
    total_sims = len(exploded_setups) * simulations_per_setup
    for this_setup in exploded_setups:
        for ss in range(simulations_per_setup):
            br = BacktestRun(parameters=this_setup, ps = PortfolioSummary(*run_backtest(**this_setup)))

            print(f'sim: {sim}/{total_sims}: setup: {[(tt, this_setup[tt]) for tt in settings_that_vary]}. Annualised return: {br.ps.portfolio_summary.annualised_ret[0]:.2f}. sharpe: {br.ps.portfolio_summary.sharpe[0]:.2f}')
            br_runs.append(br)
            sim += 1
    
    print(f'Time: {time.time()-stime}')

    ff = "../../data/portfolio_simulations_lower_min.pickle"
    pickle.dump( br_runs, open(ff, "wb"))

    # df = convert_to_ps_summary_df(br_runs)
    # df['cc'] = df['capital'].astype('category')
    # df.groupby('cc').apply(lambda x: print(x.describe()))

def all_runs_to_df(br_runs):
    return pd.concat(map(br_to_summary, br_runs), ignore_index=True)

def br_to_summary(br):
    return pd.concat([ pd.DataFrame(br.parameters, index=[0]) , br.ps.portfolio_summary], axis=1) 


def order_by_ann_return(br_runs):
    br_runs[0].ps.portfolio_summary.annualised_ret[0]



def convert_to_ps_summary_df(runs):
    bb = pd.concat(list(map(convert_df, runs)), ignore_index=True)
    return bb

def load_pickled_summary(ff):
    ff = "../../data/portfolio_simulations.pickle"
    runs_per_cap = pickle.load( open(ff, 'rb'))
    all_summaries = convert_to_summary_df(runs_per_cap)


def convert_to_summary_df(runs):
    aa = list(itertools.chain.from_iterable(runs.values()))

    bb = pd.concat(list(map(convert_df, aa)), ignore_index=True)
    return bb


def convert_df(br):
    df = br.ps.portfolio_summary
    df['capital'] = br.parameters['capital']
    return df

    

def setup_and_run():
    # logging.basicConfig(format='[ %(levelname)s ] %(asctime)s -  %(message)s', level=logging.DEBUG)
    logging.basicConfig(format='%(message)s', level=logging.INFO)

    ps, ps_summary = run_backtest(
            payment_per_stream = 0.005
            ,
            artist_advance_fraction = 0.8
            ,
            pop_5_slippage = 0.15
            , # Applied on both enter and exit
            execution_slippage = 0.1
            , # Applied only on the exit, ie how long does it take to sign up the artist
            clip_individual_returns = 10
            ,
            min_price = 5e3
            ,
            max_price = 10e3
            ,
            duration_days = 365
            ,
            new_streams_per_date = 3
            ,
            capital = 500e3
            )


def run_backtest(payment_per_stream, 
        artist_advance_fraction, 
        pop_5_slippage,
        execution_slippage, 
        clip_individual_returns, 
        min_price, 
        max_price, 
        duration_days, 
        new_streams_per_date, 
        capital):


    model_data = get_model_data()

    stream_prices = (generate_prices_for(model_data,
            artist_advance_fraction=artist_advance_fraction,
            payment_per_stream=payment_per_stream,
            pop_5_slippage=pop_5_slippage,
            execution_slippage=execution_slippage,
            clip_individual_returns=clip_individual_returns))

    stream_prices = stream_prices >> mask(X.open_price >= min_price, X.open_price <= max_price)

    randomly_generated_events = generate_random_draws(stream_prices,
            new_streams_per_date=new_streams_per_date,
            duration_days=duration_days)

    ps = execute_backtest(randomly_generated_events, capital) >> arrange('Date')

    ps['Portfolio_return'] = ps['Portfolio_worth'].pct_change()
    ps['Cumulative_portfolio_return'] = ps['Portfolio_worth']/ps['Portfolio_worth'][0] - 1

    portfolio_summary_stats = calculate_portfolio_summary_stats(ps)

    return ps, portfolio_summary_stats

def chart_portfolio_stats(ps):
    #ggplot(randomly_generated_events, aes('open_date')) + geom_bar()
    print(ggplot(ps, aes('Date','Positions')) +geom_col())
    print(ggplot(ps, aes('Date','Portfolio_return')) +geom_col())
    print(ggplot(ps, aes('Date','Cumulative_portfolio_return')) +geom_line())
    print(ggplot(ps, aes('Date','Cash')) +geom_col())


def calculate_portfolio_summary_stats(ps):
    tot_stats = ps >> summarize(average_positions = mean(X.Positions), 
                average_position_size = mean(X.Position_Size),
                average_portfolio_return = mean(X.Portfolio_return),
                portfolio_std = sd(X.Portfolio_return),
                cumulative_return = last(X.Cumulative_portfolio_return),
                total_number_tracks = ply_sum(X.entered_positions))

    rr = ps['Portfolio_return']

    tot_stats['Backtest_days'] = (ps.Date.max() - ps.Date.min())/pd.Timedelta('1 day')

    tot_stats['max_drawdown']  = pf.timeseries.max_drawdown(rr)
    tot_stats['sharpe'] = pf.timeseries.sharpe_ratio(rr)
    tot_stats['annualised_ret'] = pf.timeseries.annual_return(rr)
    tot_stats['annual_volatility']= pf.timeseries.annual_volatility(rr)
    tot_stats['sortino'] = pf.timeseries.sortino_ratio(rr)
    tot_stats['downside_risk'] = pf.timeseries.downside_risk(rr)
    tot_stats['tail_risk'] = pf.timeseries.tail_ratio(rr)

    return tot_stats

@make_symbolic
def ply_sum(x):
    return np.sum(x)

def execute_backtest(pe, capital):

    open_dates = pe >> select( X.open_date) >> distinct() >> arrange('open_date') >> pull('open_date')
    close_dates = pe >> select( X.close_date) >> distinct() >> arrange('close_date') >> pull('close_date')
    # dates = open_dates.append(close_dates, ignore_index = True).unique()
    dates = pd.date_range(start=open_dates.min(), end=close_dates.max())

    holdings = pd.DataFrame()
    trades = pd.DataFrame()
    remaining_capital = capital

    portfolio_stats = pd.DataFrame()
    i = 1

    
    for dd in dates:

        closes = (pe >> mask(X.close_date == dd))

        if len(closes) > 0:
            trades, holdings, remaining_capital = process_exits(closes, trades, holdings, remaining_capital)

        opens_by_size = (pe >> mask(X.open_date == dd) >> arrange(X.open_price))
        total_to_invest_today = (opens_by_size >> pull('open_price')).sum()
        not_entered = []

        if remaining_capital > total_to_invest_today:
            if len(opens_by_size) > 0:
                holdings, trades, remaining_capital = enter_position(opens_by_size, holdings, trades, remaining_capital)
        else:
            logging.info('Not sufficient remaining capital remaining for all positions. Entering X first')
            not_entered = opens_by_size >> mutate( cumulative_cost = cumsum(X.open_price)) >> mask(X.cumulative_cost > remaining_capital)
            opens_by_size= opens_by_size >> mutate( cumulative_cost = cumsum(X.open_price)) >> mask(X.cumulative_cost <= remaining_capital)
            logging.info(f'Skipping {len(not_entered)} position due to unsufficient remaining_capital')
            if len(opens_by_size > 0):
                holdings, trades, remaining_capital = enter_position(opens_by_size, holdings, trades, remaining_capital)

        if (len(closes) > 0) or (len(opens_by_size) > 0):

            
            logging.info(f'''
==========================================
Summary for date:   {dd} (day: {i})
Portfolio holdings: {holdings.position.sum():.0f} in {len(holdings)} positions
Remaing capital:    {remaining_capital:.0f}
Total:              {holdings.position.sum() + remaining_capital:.0f}
==========================================
            ''')


        portfolio_stats = (portfolio_stats >>
                bind_rows(
                    pd.DataFrame({
                    'Date': dd, 
                    'day' : i,
                    'Positions': len(holdings), 
                    'Position_Size': holdings.position.sum(),
                    'Cash': remaining_capital, 
                    'Portfolio_worth': holdings.position.sum() + remaining_capital,
                    'entered_positions': len(opens_by_size),
                    'Skipped tracks' : len(not_entered)}, index=[dd])))
    
        i = i + 1


    return portfolio_stats


def process_exits(closes, trades, holdings, remaining_capital):
# TODO: this might have to be the exit position price depending on what we want todo
    sell_revenue = (closes >> select(['spyid', 'close_date', 'close_price']) >>
            rename( position = 'close_price') >>
            inner_join(trades >> select('spyid'), by = 'spyid'))
    exit_trades = (sell_revenue >> select(['spyid', 'close_date']) >> inner_join(trades, on = 'spyid') >>
            select( ['spyid', 'close_date', 'position']) >> rename( date = X.close_date))

    if len(exit_trades) > 0:
        trades = trades >> bind_rows(exit_trades)
        holdings = holdings >> anti_join(closes >> select(X.spyid), by = 'spyid')
        remaining_capital = remaining_capital + sell_revenue.position.sum()

        logging.info(f'''Exiting {len(exit_trades)} positions for total revenue: {sell_revenue.position.sum():.2f} (cost of entry: {exit_trades.position.sum():.2f}).
Return on these positions: {100*(sell_revenue.position.sum() - exit_trades.position.sum())/exit_trades.position.sum():.2f}%''')

    return trades, holdings, remaining_capital

def enter_position(opens_by_size, holdings, trades, remaining_capital):
    holdings = (holdings >> bind_rows(opens_by_size >> select( ['spyid', 'open_price']) >> rename(position=X.open_price)))
    new_trades = opens_by_size >> select(['spyid', 'open_date', 'open_price']) >> rename(date=X.open_date, position=X.open_price)
    actual_investments = new_trades.position.sum()
    trades = trades >> bind_rows(new_trades)
    remaining_capital = remaining_capital - actual_investments

    logging.info(f' Entering {len(new_trades)} positions for the total position size of {new_trades["position"].sum():.0f}')

    return holdings, trades, remaining_capital




def generate_random_draws(pe, new_streams_per_date, duration_days, seed = 1):
    pretend_start_date = '2017-01-01'
    n = len(pe)
    dd = pd.date_range(pretend_start_date, periods=duration_days)
    # np.random.seed(seed)
    
    total_number = round(duration_days * new_streams_per_date)
    new_pe = pe >> sample(n = total_number, replace = True)
    new_open = np.random.choice(dd, total_number)
    new_pe = (new_pe >> 
            mutate( open_date = new_open) >> 
            mutate( close_date = X.open_date + pd.Timedelta('100 days'), spyid = np.arange(len(new_pe)) + 1) >>
            arrange('open_date'))

    return new_pe


def general_stats_base():
    model_data = pd.read_feather(
        '../../data/backtesting/pop_5_base_model.feather')

    pe = generate_prices_for(model_data)

    capital = 500e3
    minimal_price = 2e3
    unit = 1  # We always buy the same number of units

    eligible = pe >> mask(X.open_price >= minimal_price) >> arrange(
        X.open_price)
    eligible >> summarize(xx=X.open_price.sum())

    return eligible


def generate_prices_for(model_data, artist_advance_fraction, payment_per_stream, pop_5_slippage, execution_slippage, clip_individual_returns):
    """
    params:
    artist_advance_fraction: artist fraction
    payment_per_stream: payment per share
    """

    pe = (model_data >> 
            distinct(X.spyid) >> 
            mutate(open_price=X.predicted_streams * artist_advance_fraction * payment_per_stream * (1-pop_5_slippage),
                close_price=X.streams * payment_per_stream * (1 - pop_5_slippage) * (1 - execution_slippage)) >> 
            mutate(ret=(X.close_price - X.open_price) / X.open_price))

    pe['close_price'] = np.where(pe['ret'] >= clip_individual_returns, 
            pe['open_price']*(clip_individual_returns -1), 
            pe['close_price'])

    return pe >> select(['spyid','open_price', 'close_price', 'ret'])


@make_symbolic
def exp(x):
    return np.exp(x)


def general_stats_first_seen(slippage=1):
    model_data = pd.read_feather(
        '../../data/backtesting/first_seen_dummy.feather')
    af = 0.8
    pps = 5e-3

    pe = (model_data >> mutate(
        open_price=slippage * exp(X.predicted_streams) * af * pps,
        close_price=slippage * exp(X.actual_streams) * pps) >>
          mutate(ret=(X.close_price - X.open_price) / X.open_price))

    capital = 500e3
    minimal_price = 2e3
    unit = 1  # We always buy the same number of units

    eligible = pe >> mask(X.open_price >= minimal_price) >> arrange(
        X.open_price)
    eligible >> summarize(xx=X.open_price.sum())

    return eligible