from progpy.datasets import nasa_battery
(desc, data) = nasa_battery.load_data(1)

print(desc['description'])

data[35]['absoluteTime'] = data[35]['relativeTime']
for i in range(36, 50):
    data[i]['absoluteTime'] = data[i]['relativeTime'] + data[i-1]['absoluteTime'].iloc[-1]

random_walk_dataset = pd.concat(data[35:50], ignore_index=True)
print(random_walk_dataset)
random_walk_dataset.plot(y=['current', 'voltage', 'temperature'], subplots=True, xlabel='Time (sec)')

from progpy.models import BatteryElectroChemEOD
batt: BatteryElectroChemEOD = BatteryElectroChemEOD()

batt['Ro'] = 0.15
batt['qMobile'] = 7750

import numpy as np
from progpy.state_estimators import ParticleFilter
from progpy.uncertain_data import MultivariateNormalDist

initial_state = batt.initialize() # Initialize model
# Define distribution around initial state
x_guess = MultivariateNormalDist(
    labels=initial_state.keys(),
    mean=initial_state.values(),
    covar=np.diag([max(1e-9, abs(x)) for x in initial_state.values()])
)

pf = ParticleFilter(batt, x_guess)

from progpy.predictors import MonteCarlo

PROCESS_NOISE = 2e-4           # Percentage process noise
MEASUREMENT_NOISE = 1e-4        # Percentage measurement noise

# Apply process noise to state
batt.parameters['process_noise'] = {key: PROCESS_NOISE * value for key, value in initial_state.items()}

# Apply measurement noise to output
z0 = batt.output(initial_state)
batt.parameters['measurement_noise'] = {key: MEASUREMENT_NOISE * value for key, value in z0.items()}

mc = MonteCarlo(batt)

NUM_SAMPLES = 100
STEP_SIZE = 1
PREDICTION_HORIZON = random_walk_dataset['absoluteTime'].iloc[-1] 

# Extract time and outputs from data
times_rw = random_walk_dataset['absoluteTime']
outputs_rw = [{'v': elem[1]['voltage']} for elem in random_walk_dataset.iterrows()]

# Define function
import numpy as np
def future_load_rw(t, x=None):
    current = np.interp(t, times_rw, random_walk_dataset['current'])
    return {'i': current}

batt.parameters['VEOD'] = 3.3

PREDICTION_UPDATE_FREQ = 50

from progpy.predictors import ToEPredictionProfile
profile = ToEPredictionProfile()

# Loop through time
for ind in range(3, random_walk_dataset.shape[0]):
    # Extract data
    t = random_walk_dataset['absoluteTime'][ind]
    i = {'i': random_walk_dataset['current'][ind]}
    z = {'t': random_walk_dataset['temperature'][ind], 'v': random_walk_dataset['voltage'][ind]}

    # Perform state estimation 
    pf.estimate(t, i, z)
    eod = batt.event_state(pf.x.mean)['EOD']
    print("  - Event State: ", eod)

    # Prediction step (at specified frequency)
    if (ind%PREDICTION_UPDATE_FREQ == 0):
        # Perform prediction
        mc_results = mc.predict(pf.x, future_load_rw, t0 = t, n_samples=NUM_SAMPLES, dt=1, horizon=PREDICTION_HORIZON, const_load=True)
        
        # Calculate metrics and print
        metrics = mc_results.time_of_event.metrics()
        print('  - ToE: {} (sigma: {})'.format(metrics['EOD']['mean'], metrics['EOD']['std']))

        # Save results
        profile.add_prediction(t, mc_results.time_of_event)

cra = profile.cumulative_relative_accuracy(GROUND_TRUTH)
print(f"Cumulative Relative Accuracy for 'EOD': {cra['EOD']}")

GROUND_TRUTH= {'EOD': 1600} 
ALPHA = 0.05
playback_plots = profile.plot(GROUND_TRUTH, ALPHA, True)