import pandas as pd
import numpy as np
from random import normalvariate, random
import plotly.express as px

from radcad import Model, Simulation, Experiment
import streamlit as st


# Additional dependencies

# For analytics
import numpy as np
# For visualization
import plotly.express as px

from PIL import Image
# Additional dependencies


pd.options.plotting.backend = "plotly"

st.header('DeSci Value Flow Model')
image = Image.open('desci.png')
st.image(image, caption='DeSci value flow schema')


def p_researcher1(params, substep, state_history, previous_state):
    losses = 0
    to_market = 0
    to_researcher = 0
    to_treasury = 0
    to_other_researcher = 0
    salary = 0
    funding = 0
    if (previous_state['timestep'] < params['timestep_switch']) and (previous_state['funding_pool'] > params['funding_round']):
        funding = params['funding_round']
        to_treasury -= funding
        research_value = funding * (1-params['epsilon'])
        losses += funding - research_value

        salary = research_value * params['beta']
        to_market = research_value
        if (random() < params['probability_buying']) and (previous_state['researcher1_value'] > params['cost_buying']):
            salary = salary - params['cost_buying']
            tx_fee = params['cost_buying'] * params['tx_fee']
            salary -= tx_fee
            to_treasury += tx_fee
            to_other_researcher += params['cost_buying']
        to_researcher += salary + research_value
    elif (previous_state['timestep'] > params['timestep_switch']) and (previous_state['researcher1_value'] > params['cost_buying']):
        tx_fee = params['cost_buying'] * params['tx_fee']
        to_researcher -= params['cost_buying'] - tx_fee
        to_other_researcher += params['cost_buying']
        to_treasury += tx_fee
    return {'update_researcher1_funding': funding,
                'update_researcher1_salary': salary,
                'update_researcher1_value': to_researcher,
                'update_funding_pool': to_treasury,
                'update_market': to_market,
                'update_researcher2_value': to_other_researcher,
                'update_losses': losses}

def p_researcher2(params, substep, state_history, previous_state):
    losses = 0
    to_market = 0
    to_researcher = 0
    to_treasury = 0
    to_other_researcher = 0
    salary = 0
    funding = 0
    if (previous_state['timestep'] > params['timestep_switch']) and (previous_state['funding_pool'] > params['funding_round']):
        funding = params['funding_round']
        to_treasury -= funding
        research_value = funding * (1-params['epsilon'])
        losses += funding - research_value

        salary = research_value * params['beta']
        to_market = research_value
        if (random() < params['probability_buying']) and (previous_state['researcher2_value'] > params['cost_buying']):
            salary = salary - params['cost_buying']
            tx_fee = params['cost_buying'] * params['tx_fee']
            salary -= tx_fee
            to_treasury += tx_fee
            to_other_researcher += params['cost_buying']
        to_researcher += salary + research_value
    elif (previous_state['timestep'] < params['timestep_switch']) and (previous_state['researcher2_value'] > params['cost_buying']):
        tx_fee = params['cost_buying'] * params['tx_fee']
        to_researcher -= params['cost_buying'] - tx_fee
        to_other_researcher += params['cost_buying']
        to_treasury += tx_fee
    return {'update_researcher2_funding': funding,
                'update_researcher2_salary': salary,
                'update_researcher2_value': to_researcher,
                'update_funding_pool': to_treasury,
                'update_market': to_market,
                'update_researcher1_value': to_other_researcher,
                'update_losses': losses}

def s_timestep(params, substep, state_history, previous_state, policy_input):
    updated_timestep = previous_state['timestep'] + 1
    return 'timestep', updated_timestep

def s_funding_pool(params, substep, state_history, previous_state, policy_input):
    funding_pool = previous_state['funding_pool']
    updated_funding_pool = funding_pool + policy_input['update_funding_pool']
    return 'funding_pool', updated_funding_pool

def s_researcher1_value(params, substep, state_history, previous_state, policy_input):
    r_value = previous_state['researcher1_value']
    updated_researcher1_value = r_value + policy_input['update_researcher1_value']
    return 'researcher1_value', updated_researcher1_value

def s_researcher1_funding(params, substep, state_history, previous_state, policy_input):
    r_funding = previous_state['researcher1_funding']
    updated_researcher1_funding = r_funding + policy_input['update_researcher1_funding']
    return 'researcher1_funding', updated_researcher1_funding

def s_researcher1_salary(params, substep, state_history, previous_state, policy_input):
    r_salary = previous_state['researcher1_salary']
    updated_researcher1_salary = r_salary + policy_input['update_researcher1_salary']
    return 'researcher1_salary', updated_researcher1_salary

def s_researcher2_value(params, substep, state_history, previous_state, policy_input):
    r_value = previous_state['researcher2_value']
    updated_researcher2_value = r_value + policy_input['update_researcher2_value']
    return 'researcher2_value', updated_researcher2_value

def s_researcher2_funding(params, substep, state_history, previous_state, policy_input):
    r_funding = previous_state['researcher2_funding']
    updated_researcher2_funding = r_funding + policy_input['update_researcher2_funding']
    return 'researcher2_funding', updated_researcher2_funding

def s_researcher2_salary(params, substep, state_history, previous_state, policy_input):
    r_salary = previous_state['researcher2_salary']
    updated_researcher2_salary = r_salary + policy_input['update_researcher2_salary']
    return 'researcher2_salary', updated_researcher2_salary

def s_knowledge_market(params, substep, state_history, previous_state, policy_input):
    value = previous_state['knowledge_market_value']
    updated_market_value = value + policy_input['update_market']
    return 'knowledge_market_value', updated_market_value

def s_losses(params, substep, state_history, previous_state, policy_input):
    losses = previous_state['losses']
    updated_losses = losses + policy_input['update_losses']
    return 'losses', updated_losses

st.subheader('Initial Value Allocation')
funding_pool = st.slider('Initial Funding Pool', min_value=1000, max_value=10000, value=1000, step=10)
researcher1_value = st.slider('Researcher1 Tokens', min_value=0, max_value=1000, value=0, step=1)
researcher2_value = st.slider('Researcher2 Tokens', min_value=0, max_value=1000, value=0, step=1)
st.subheader('Simulation Parameters')
tx_fee = st.slider('Transaction fee collected by DAO treasury during each transaction in the knowledge market', min_value=0., max_value=1., value=0.1, step=0.0001)
st.write('Set the funding disbursed each round from the funding pool')
funding_round = st.slider('Funding Round', min_value=100, max_value=1000, value=100, step=1)
st.write('Set the relative value leakages in the model.')
epsilon = st.slider('Work Inefficiency Weight', min_value=0., max_value=1., value=0.1, step=0.0001)
st.write('Set the portion of grant funding to be used as researcher salary.')
beta = st.slider('Salary Weight', min_value=0., max_value=1., value=0.4, step=0.0001)
st.write('Set the cost of getting access to papers in the knowledge market.')
cost_buying = st.slider('Cost of Buying', min_value=10., max_value=100., value=10., step=0.1)
st.write('Set the probability a researcher will buy access to a paper at each timestep.')
probability_buying = st.slider('Researcher Probability of Buying', min_value=0., max_value=1., value=0.1, step=0.0001)
st.write('Set the number of timesteps in the simulation.')
timesteps = st.slider('Timesteps', min_value=10, max_value=1000, value=100, step=1)

initial_state = {
    'funding_pool': funding_pool,
    'researcher1_value': researcher1_value,
    'researcher1_funding': 0,
    'researcher1_salary': 0,
    'researcher2_value': researcher2_value,
    'researcher2_funding': 0,
    'researcher2_salary': 0,
    'knowledge_market_value': 0,
    'timestep': 0,
    'losses': 0
}
ts = int(timesteps/2)

system_params = {
    'funding_pool': [funding_pool],
    'funding_round': [funding_round],
    'beta': [beta],
    'epsilon': [epsilon],
    'cost_buying': [cost_buying],
    'probability_buying': [probability_buying],
    'timestep_switch': [ts],
    'tx_fee': [tx_fee]
}

def configure_and_run_experiment(initial_state,
                      partial_state_update_blocks,
                      timesteps):
    model = Model(
    # Model initial state
    initial_state=initial_state,
    # Model Partial State Update Blocks
    state_update_blocks=partial_state_update_blocks,
    # System Parameters
    params=system_params
    )
    simulation = Simulation(
    model=model,
    timesteps=timesteps,  # Number of timesteps
    runs=1  # Number of Monte Carlo Runs
    )   

    result = simulation.run()
    return result

partial_state_update_blocks = [
    {
        'policies': {
            'p_researcher1': p_researcher1,
            'p_researcher2': p_researcher2
        },
        'variables': {
            'timestep': s_timestep,
            'funding_pool': s_funding_pool,
            'researcher1_value': s_researcher1_value,
            'researcher1_funding': s_researcher1_funding,
            'researcher1_salary': s_researcher1_salary,
            'researcher2_value': s_researcher2_value,
            'researcher2_funding': s_researcher2_funding,
            'researcher2_salary': s_researcher2_salary,
            'knowledge_market_value': s_knowledge_market,
            'losses': s_losses
        }
    }
]

if st.button('Run Simulation'):
    raw_result = configure_and_run_experiment(initial_state, partial_state_update_blocks, timesteps)
    df = pd.DataFrame(raw_result)
    fig1 = df.plot(kind='line', x='timestep', y=['funding_pool', 'researcher1_value', 'researcher2_value'], width=1000)
    fig2 = df.plot(kind='line', x='timestep', y=['funding_pool','knowledge_market_value'], width=1000)
    fig3 = df.plot(kind='line', x='timestep', y=['funding_pool', 'losses'], width=1000)
    fig4 = df.plot(kind='line', x='timestep', y=['researcher1_value', 'researcher2_value', 'losses'], width=1000)
    st.subheader('Results')
    st.plotly_chart(fig1)
    st.plotly_chart(fig2)
    st.plotly_chart(fig3)
    st.plotly_chart(fig4)