# =====================================================
# utils.py
# Backend Functions for Portfolio Allocation
# =====================================================

import os
import joblib
import numpy as np
import pandas as pd

from stable_baselines3 import PPO

from config import (
    PPO_PATH,
    RL_FEATURE_DATA,
    PRICE_DATA,
    CORRELATION_MATRIX,
    ALL_STOCKS,
    RISK_MAPPING
)

# =====================================================
# Load PPO Model
# =====================================================

def load_ppo():

    model = PPO.load(PPO_PATH)

    return model

# =====================================================
# Load Deployment Data
# =====================================================

def load_data():

    rl_feature_data = joblib.load(RL_FEATURE_DATA)

    price_data = joblib.load(PRICE_DATA)

    correlation_matrix = joblib.load(CORRELATION_MATRIX)

    return (

        rl_feature_data,

        price_data,

        correlation_matrix

    )

# =====================================================
# Latest Feature Vector
# =====================================================

def get_latest_feature_matrix(
    rl_feature_data
):

    feature_matrix = []

    for stock in ALL_STOCKS:

        latest = (

            rl_feature_data[stock]

            .iloc[-1]

            .values

            .astype(np.float32)

        )

        feature_matrix.extend(latest)

    return np.array(
        feature_matrix,
        dtype=np.float32
    )

def build_observation(

    rl_feature_data,

    budget,

    risk_profile,

    investment_horizon

):

    feature_matrix = get_latest_feature_matrix(
        rl_feature_data
    )

    shares = np.zeros(
        len(ALL_STOCKS),
        dtype=np.float32
    )

    cash = budget

    portfolio_value = budget

    if budget < 100000:

        tier = 0

    elif budget < 500000:

        tier = 1

    else:

        tier = 2

    portfolio_state = np.array(

        [

            budget,

            tier,

            cash,

            RISK_MAPPING[risk_profile],

            investment_horizon,

            portfolio_value

        ],

        dtype=np.float32

    )

    observation = np.concatenate(

        [

            feature_matrix,

            shares,

            portfolio_state

        ]

    )

    return observation

def predict_portfolio(

    model,

    observation

):

    action, _ = model.predict(

        observation,

        deterministic=True

    )

    action = np.clip(

        action,

        0,

        None

    )

    if action.sum() == 0:

        action += 1

    weights = action / action.sum()

    return weights

def generate_portfolio(

    weights,

    budget

):

    allocation = pd.DataFrame({

        "Stock": ALL_STOCKS,

        "Weight (%)": weights[:-1] * 100,

        "Investment (₹)": weights[:-1] * budget

    })

    allocation = allocation.sort_values(

        "Weight (%)",

        ascending=False

    )

    allocation.reset_index(

        drop=True,

        inplace=True

    )

    return allocation

def cash_remaining(

    weights,

    budget

):

    return weights[-1] * budget

def portfolio_summary(

    weights,

    budget

):

    return {

        "Total Investment":

            budget -

            cash_remaining(

                weights,

                budget

            ),

        "Cash":

            cash_remaining(

                weights,

                budget

            ),

        "Number of Stocks":

            np.sum(

                weights[:-1] > 0

            )

    }

# =====================================================
# Main Recommendation Pipeline
# =====================================================

def generate_recommendation(
    budget,
    risk_profile,
    investment_horizon
):

    # Load everything
    model = load_ppo()

    rl_feature_data, price_data, correlation_matrix = load_data()

    snapshot_date = (
    pd.to_datetime(
        next(iter(rl_feature_data.values())).index[-1]
    ).strftime("%d-%b-%Y"))

    # Build PPO observation
    observation = build_observation(
        rl_feature_data=rl_feature_data,
        budget=budget,
        risk_profile=risk_profile,
        investment_horizon=investment_horizon
    )

    # PPO prediction
    weights = predict_portfolio(
        model,
        observation
    )

    # Portfolio table
    allocation = generate_portfolio(
        weights,
        budget
    )

    # Cash
    cash = cash_remaining(
        weights,
        budget
    )

    # Summary
    summary = portfolio_summary(
        weights,
        budget
    )

    return allocation, cash, summary, weights, snapshot_date