from steel_processes import define_processes
processes = define_processes()

import theano
theano.config.mode = 'FAST_RUN'
theano.config.optimizer = 'fast_compile'

input_defs = {
    'BF': 5000,
    'DR': 300,
    'SP': 1000,
    'IFC': 300,
}

import pickle
from priors import param_defs
from leontief_model import SplitParamModel

model1 = SplitParamModel(processes, input_defs, param_defs)

with open('model1.pickle', 'wb') as f:
    pickle.dump(model1, f, protocol=pickle.HIGHEST_PROTOCOL)

import pymc3 as pm
RANDOM_SEED = 974067  # from random.org

with model1.model:
    advi1 = pm.variational.advi(n=10000, random_seed=RANDOM_SEED)
np.savez('traces/advi1.npz', advi1)

plt.plot(advi1.elbo_vals[0:]);

def sample(model, advi, trace=None, num_samples=2000):
    step = pm.NUTS(scaling=model.dict_to_array(advi.stds)**2, is_cov=True)
    trace = pm.sample(num_samples, start=advi.means, step=step, trace=trace)
    return trace

with model1.model:
    trace1 = sample(model1.model, advi1, trace=pm.backends.Text('traces/trace1'))

pm.traceplot(trace1, varnames=['inputs', 'param_OBF', 'param_OBFS']);

observations1 = [
        (['BF'], ['PI'], 928.4, 50),  # worldsteel: Table 41 - total pig iron
        (['DR'], ['DRI'], 65.8, 6),   # worldsteel: Table 42 - total DRI
        
        (['EAF'], ['EAFS'], 407.0, 40),  # worldsteel: Table 6 - EAF crude steel production
        (['OBF'], ['OBFS'], 891.7, 80),  # worldsteel: Table 6 - OBC crude steel production
        (['OHF'], ['OHFS'], 29.3, 3),    # worldsteel: Table 6 - OHF + other crude steel production
        
        # SEM output, from worldsteel
        (['SEMP'], ['seclight'], 44.4, 2),    # Cullen flow [105]: worldsteel table 16
        (['SEMP'], ['secheavy'], 39.6, 2),    # Cullen flow [106]: worldsteel table 15
        (['SEMP'], ['secrail'], 10.0, 1),     # Cullen flow [107]: worldsteel table 14
        
        # RBM output, from worldsteel
        (['STP'], ['seamlesstube'], 27.2, 2), # Cullen flow [85]: worldsteel table 25
        (['RBMP'], ['rodrebar'], 147.0, 15),  # Cullen flow [108]: worldsteel table 17 (he added extra 26.6 to make it add up)
        (['RBMP'], ['rodwire'], 148.5, 15),   # Cullen flow [109]: worldsteel table 19
        (['RBMP'], ['rodbar'], 109.7, 11),    # Cullen flow [110]: worldsteel table 18
        
        # PLM output, from worldsteel & others
        (['TWP'], ['weldedtube'], 44.5, 4),   # Cullen flow [89]: worldsteel table 26. (he solved as 62.4 to make "apparent consumption" match)
        (['PLMP'], ['plate'], 110.0, 11),     # Cullen flow [111]: SBB value (ref 14, p4)
        
        # HSM output, from worldsteel & SBB
        (['CRMP'], ['elsheet'], 10.3, 2),      # Cullen flow [112]: worldsteel table 20
        (['CRMP'], ['crc'], 135.0, 13),        # Cullen flow [113]: SBB (ref 14, p4)
        (['GPP_CR'], ['crcgalv'], 96.4, 9),    # Cullen flow [114]: worldsteel table 22 (assumed all CRC)
        (['HSMP'], ['hrc'], 150.0, 15),        # Cullen flow [115]: SBB (ref 14, p4)
        (['HSMP'], ['hrns'], 40.0, 4),         # Cullen flow [116]: SBB (ref 14, p4)
        (['GP_HR'], ['hrcgalv'], 10.0, 2),     # Cullen flow [97]: SBB (ref 14, p4)
        (['TM'], ['crctinned'], 11.6, 2),      # Cullen flow [100]: worldsteel table 21 (all tinmill)
        (['OCP'], ['crccoated'], 16.5, 2),     # Cullen flow [103]: worldsteel table 23 (all non-metalic coated)
]

model2 = SplitParamModel(processes, input_defs, param_defs, flow_observations=observations1)

with open('model2.pickle', 'wb') as f:
    pickle.dump(model2, f, protocol=pickle.HIGHEST_PROTOCOL)

with model2.model:
    advi2 = pm.variational.advi(n=10000)
np.savez('traces/advi2.npz', advi2)    

plt.plot(advi2.elbo_vals[0:]);

with model2.model:
    trace2 = sample(model2.model, advi2, trace=pm.backends.Text('traces/trace2'), num_samples=500)

pm.traceplot(trace2, varnames=['inputs', 'param_OBF', 'param_OBFS', 'Fobs']);

pm.traceplot(trace2, varnames=['inputs', 'param_OBF', 'param_OBFS', 'Fobs']);

observations2 = observations1 + [
    (['PI'], ['EAF'], 44.6, 4),   # Cullen flow [12] (ref 2)
    (['SPC'], ['caststeel'], 10.5, 2),    # Cullen flow [51], from ref (11)
    (['IFC'], ['castiron'], 68.3, 6),     # Cullen flow [57], from ref (11)
]

input_observations = [
    (['SP'], 475.5, 40),     # Cullen flow [10]; change if fabrication scrap is added
]

inflow_fraction_observations = [
    (['S'], ['IFC'], 0.17, 0.03),         # Cullen flow [54] -- scrap is 17%.
    # Should be relative to cast iron output, not process throughput.
]

model3 = SplitParamModel(processes, input_defs, param_defs,
                         flow_observations=observations2,
                         input_observations=input_observations,
                         inflow_observations=inflow_fraction_observations)

with open('model3.pickle', 'wb') as f:
    pickle.dump(model3, f, protocol=pickle.HIGHEST_PROTOCOL)

with model3.model:
    advi3 = pm.variational.advi(n=10000)
np.savez('traces/advi3.npz', advi3)

plt.plot(advi3.elbo_vals[0:]);

with model3.model:
    trace3 = sample(model3.model, advi3, trace=pm.backends.Text('traces/trace3'), num_samples=500)

pm.traceplot(trace3, varnames=['inputs', 'param_OBF', 'param_OBFS', 'Fobs']);