import os
import numpy as np
import pandas as pd
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from collections import OrderedDict
from tensorflow.python.keras.models import load_model
from pkg_resources import resource_filename

from transomaly.prepare_input import PrepareInputArrays
from transomaly.loss_functions import mean_squared_error, chisquare_loss, mean_squared_error_over_error
matplotlib.use('TkAgg')

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
COLPB = {'g': 'tab:green', 'r': 'tab:red'}
MARKPB = {'g': 'o', 'r': 's', 'z': 'd'}
ALPHAPB = {'g': 0.3, 'r': 1., 'z': 1}

CLASS_COLOR = {'SNIa-norm': 'tab:green', 'SNIbc': 'tab:orange', 'SNII': 'tab:blue', 'SNIIn': 'blue',
               'SNIa-91bg': 'tab:red', 'SNIa-x': 'bisque', 'point-Ia': 'tab:brown', 'Kilonova': '#aaffc3',
               'SLSN-I': 'tab:olive', 'PISN': 'tab:cyan', 'ILOT': '#FF1493', 'CART': 'navy', 'TDE': 'tab:pink',
               'AGN': 'tab:purple'}

# npred = 7
# model_filepath_onepoint_infuture = "/Users/danmuth/OneDrive - University of Cambridge/PycharmProjects/transomaly/plots/model__ci()_ns1_c(1,)/keras_model_epochs300_onepoint_pred7timesteps_infuture_normalised_predict_last49_timesteps_nodropout_100lstmneurons/keras_model_epochs300_onepoint_pred7timesteps_infuture_normalised_predict_last49_timesteps_nodropout_100lstmneurons.hdf5"
model_filepath = "/Users/danmuth/OneDrive - University of Cambridge/PycharmProjects/transomaly/plots/model__ci()_ns1_c(1,)/keras_model_epochs500_pred2timesteps_normalised_predict_last49_timesteps_nodropout_100lstmneurons/keras_model_epochs500_pred2timesteps_normalised_predict_last49_timesteps_nodropout_100lstmneurons.hdf5"

model = load_model(model_filepath, custom_objects={'loss': mean_squared_error()})

passbands = ('g','r')
contextual_info = ()

X = np.array([np.array([[0.09158034, 0.07176773],
       [0.09008677, 0.07137485],
       [0.08917016, 0.0727186 ],
       [0.09029362, 0.07340094],
       [0.08999084, 0.07224263],
       [0.08900606, 0.07197019],
       [0.08739904, 0.07095805],
       [0.08682939, 0.07071227],
       [0.08567506, 0.07031706],
       [0.08386102, 0.07021409],
       [0.0836062 , 0.0698832 ],
       [0.08332578, 0.07050169],
       [0.08397429, 0.06937913],
       [0.0860095 , 0.0716971 ],
       [0.09076107, 0.07780995],
       [0.09871331, 0.08934081],
       [0.11663096, 0.11012718],
       [0.15700709, 0.15327507],
       [0.23417453, 0.22872161],
       [0.35862168, 0.33963544],
       [0.51792838, 0.48456903],
       [0.68616083, 0.64401235],
       [0.81866526, 0.79203408],
       [0.90768435, 0.89154676],
       [0.93726188, 0.94250969],
       [0.92312219, 0.95238174],
       [0.85896026, 0.93542825],
       [0.74499556, 0.87812468],
       [0.62114839, 0.78786493],
       [0.50529031, 0.6932507 ],
       [0.40645818, 0.6036984 ],
       [0.32993884, 0.52479434],
       [0.27163636, 0.45807863],
       [0.22780797, 0.40140124],
       [0.19817062, 0.35339631],
       [0.17674938, 0.31370617],
       [0.16142444, 0.27973599],
       [0.14923072, 0.25118097],
       [0.13910386, 0.22674725],
       [0.13288632, 0.20717057],
       [0.12727512, 0.191151  ],
       [0.12253798, 0.17674126],
       [0.11748478, 0.16368262],
       [0.11428479, 0.15247652],
       [0.110435  , 0.14285344],
       [0.10726069, 0.13410306],
       [0.10317767, 0.12546451],
       [0.10019489, 0.11729226],
       [0.09654251, 0.11025106]])])
y = np.array([np.array([[0.09008677, 0.07137485],
       [0.08917016, 0.0727186 ],
       [0.09029362, 0.07340094],
       [0.08999084, 0.07224263],
       [0.08900606, 0.07197019],
       [0.08739904, 0.07095805],
       [0.08682939, 0.07071227],
       [0.08567506, 0.07031706],
       [0.08386102, 0.07021409],
       [0.0836062 , 0.0698832 ],
       [0.08332578, 0.07050169],
       [0.08397429, 0.06937913],
       [0.0860095 , 0.0716971 ],
       [0.09076107, 0.07780995],
       [0.09871331, 0.08934081],
       [0.11663096, 0.11012718],
       [0.15700709, 0.15327507],
       [0.23417453, 0.22872161],
       [0.35862168, 0.33963544],
       [0.51792838, 0.48456903],
       [0.68616083, 0.64401235],
       [0.81866526, 0.79203408],
       [0.90768435, 0.89154676],
       [0.93726188, 0.94250969],
       [0.92312219, 0.95238174],
       [0.85896026, 0.93542825],
       [0.74499556, 0.87812468],
       [0.62114839, 0.78786493],
       [0.50529031, 0.6932507 ],
       [0.40645818, 0.6036984 ],
       [0.32993884, 0.52479434],
       [0.27163636, 0.45807863],
       [0.22780797, 0.40140124],
       [0.19817062, 0.35339631],
       [0.17674938, 0.31370617],
       [0.16142444, 0.27973599],
       [0.14923072, 0.25118097],
       [0.13910386, 0.22674725],
       [0.13288632, 0.20717057],
       [0.12727512, 0.191151  ],
       [0.12253798, 0.17674126],
       [0.11748478, 0.16368262],
       [0.11428479, 0.15247652],
       [0.110435  , 0.14285344],
       [0.10726069, 0.13410306],
       [0.10317767, 0.12546451],
       [0.10019489, 0.11729226],
       [0.09654251, 0.11025106],
       [0.09630747, 0.10437187]])])
timesX = np.array([np.array([-70., -67., -64., -61., -58., -55., -52., -49., -46., -43., -40.,
       -37., -34., -31., -28., -25., -22., -19., -16., -13., -10.,  -7.,
        -4.,  -1.,   2.,   5.,   8.,  11.,  14.,  17.,  20.,  23.,  26.,
        29.,  32.,  35.,  38.,  41.,  44.,  47.,  50.,  53.,  56.,  59.,
        62.,  65.,  68.,  71.,  74.,  77.])])
objids = np.array(['median_Ia'])

X = np.array([np.array([[0.00448964, 0.        ],
       [0.        , 0.02469979],
       [0.00713832, 0.0510595 ],
       [0.02991802, 0.08449011],
       [0.05011106, 0.09195332],
       [0.05702283, 0.07957474],
       [0.05513788, 0.05611958],
       [0.04902122, 0.0294538 ],
       [0.04288751, 0.00993177],
       [0.04127183, 0.01818543],
       [0.05353569, 0.04212558],
       [0.06324364, 0.06195425],
       [0.05741975, 0.07571231],
       [0.03920042, 0.08415738],
       [0.02232076, 0.08873083],
       [0.01912765, 0.08509039],
       [0.03159054, 0.07017496],
       [0.04909139, 0.05580469],
       [0.06360378, 0.05487265],
       [0.08042658, 0.08175915],
       [0.14612581, 0.15545786],
       [0.32463802, 0.29112003],
       [0.5616513 , 0.4847463 ],
       [0.78721409, 0.70134517],
       [0.94289831, 0.88820695],
       [1.        , 0.98685417],
       [0.9617599 , 1.        ],
       [0.84395464, 0.95538324],
       [0.70269411, 0.87417796],
       [0.56370832, 0.77777756],
       [0.43422203, 0.6788543 ],
       [0.31954263, 0.58431753],
       [0.22572829, 0.5003608 ],
       [0.15997532, 0.43230079],
       [0.13003542, 0.3766077 ],
       [0.11756719, 0.32962512],
       [0.1001643 , 0.29821954],
       [0.08371692, 0.28562479],
       [0.07356015, 0.27645137],
       [0.06571214, 0.2572178 ],
       [0.05947788, 0.23700196],
       [0.05453582, 0.21703187],
       [0.05062684, 0.19788037],
       [0.04754242, 0.17990707],
       [0.04511494, 0.16331596],
       [0.04320991, 0.14819947],
       [0.04171951, 0.1345717 ],
       [0.04055748, 0.12239351],
       [0.03965491, 0.11159111]])])
y = np.array([np.array([[0.        , 0.02469979],
       [0.00713832, 0.0510595 ],
       [0.02991802, 0.08449011],
       [0.05011106, 0.09195332],
       [0.05702283, 0.07957474],
       [0.05513788, 0.05611958],
       [0.04902122, 0.0294538 ],
       [0.04288751, 0.00993177],
       [0.04127183, 0.01818543],
       [0.05353569, 0.04212558],
       [0.06324364, 0.06195425],
       [0.05741975, 0.07571231],
       [0.03920042, 0.08415738],
       [0.02232076, 0.08873083],
       [0.01912765, 0.08509039],
       [0.03159054, 0.07017496],
       [0.04909139, 0.05580469],
       [0.06360378, 0.05487265],
       [0.08042658, 0.08175915],
       [0.14612581, 0.15545786],
       [0.32463802, 0.29112003],
       [0.5616513 , 0.4847463 ],
       [0.78721409, 0.70134517],
       [0.94289831, 0.88820695],
       [1.        , 0.98685417],
       [0.9617599 , 1.        ],
       [0.84395464, 0.95538324],
       [0.70269411, 0.87417796],
       [0.56370832, 0.77777756],
       [0.43422203, 0.6788543 ],
       [0.31954263, 0.58431753],
       [0.22572829, 0.5003608 ],
       [0.15997532, 0.43230079],
       [0.13003542, 0.3766077 ],
       [0.11756719, 0.32962512],
       [0.1001643 , 0.29821954],
       [0.08371692, 0.28562479],
       [0.07356015, 0.27645137],
       [0.06571214, 0.2572178 ],
       [0.05947788, 0.23700196],
       [0.05453582, 0.21703187],
       [0.05062684, 0.19788037],
       [0.04754242, 0.17990707],
       [0.04511494, 0.16331596],
       [0.04320991, 0.14819947],
       [0.04171951, 0.1345717 ],
       [0.04055748, 0.12239351],
       [0.03965491, 0.11159111],
       [0.03895684, 0.10206975]])])
timesX = np.array([np.array([-70., -67., -64., -61., -58., -55., -52., -49., -46., -43., -40.,
       -37., -34., -31., -28., -25., -22., -19., -16., -13., -10.,  -7.,
        -4.,  -1.,   2.,   5.,   8.,  11.,  14.,  17.,  20.,  23.,  26.,
        29.,  32.,  35.,  38.,  41.,  44.,  47.,  50.,  53.,  56.,  59.,
        62.,  65.,  68.,  71.,  74.,  77.])])
objids = np.array(['1_6222236'])

# X[0,:,1][25:] = 0
# X[0,:,0][25:] = 0
# y[0,:,1][24:] = 0
# y[0,:,0][24:] = 0
# X[0,:,2] = 0.04

npred = 2
tidx = -npred+1 if npred != 1 else None

y_pred = model.predict(X[:,:tidx,:2])

idx = 0
sidx = 0


for p in range(50-npred):
       # y_pred_correct_format = np.zeros(y.shape)
       # y_pred_correct_format[:, :, 0] = np.copy(y_predict[:, p, :49])
       # y_pred_correct_format[:, :, 1] = np.copy(y_predict[:, p, :])

       fig, (ax1) = plt.subplots(nrows=1, ncols=1)

       for pbidx, pb in enumerate(passbands):
           plotlabeltest = "ytest:{}".format(pb)
           plotlabelpred = "ypred:{}".format(pb)
           marker = None  # MARKPB[pb] if s == 0 else None
           ax1.plot(timesX[sidx][1:], y[sidx][:, pbidx], c=COLPB[pb], lw=1,
                    label=plotlabeltest, marker=None, markersize=10, alpha=1, linestyle='-')
           ax1.plot(timesX[sidx][1:tidx], y_pred[sidx][:, pbidx+2], c=COLPB[pb], lw=1,
                    label=plotlabelpred, marker=None, markersize=10, alpha=1, linestyle=':')
           # ax1.plot(timesX[sidx][(1+p):(1+p+npred)], y_pred[sidx][p][pbidx::2], c=COLPB[pb], lw=1,
           #          label=plotlabelpred, marker=None, markersize=10, alpha=1, linestyle=':')
           ax1.axvspan(timesX[sidx][1+p], timesX[sidx][-1], alpha=0.2, color='grey')

           # pbmask = lc['passband'] == pb
           # sortedidx = np.argsort(lc[pbmask]['time'].data)
           # time = lc[pbmask]['time'].data[sortedidx]
           # flux = lc[pbmask]['flux'].data[sortedidx]
           # fluxerr = lc[pbmask]['fluxErr'].data[sortedidx]

           # ax1.errorbar(time, flux, yerr=fluxerr,
           #              fmt=".", capsize=0, color=COLPB[pb], label='_nolegend_')

       ax1.set_ylabel("Relative flux")
       ax1.set_xlabel("Time since trigger [days]")
       ax1.legend()
       figdir = 'interesting plots examining what RNN has learned/predict_npred_timesteps_ateachstep/'
       plt.savefig(figdir+f'predict-{npred}-timesteps-{objids[0]}_ateachstep_{p}.png')
       # plt.show()







import os
import numpy as np
import pandas as pd
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from collections import OrderedDict
from tensorflow.python.keras.models import load_model
from pkg_resources import resource_filename

from transomaly.prepare_input import PrepareInputArrays
from transomaly.loss_functions import mean_squared_error, chisquare_loss, mean_squared_error_over_error
matplotlib.use('TkAgg')

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
COLPB = {'g': 'tab:green', 'r': 'tab:red'}
MARKPB = {'g': 'o', 'r': 's', 'z': 'd'}
ALPHAPB = {'g': 0.3, 'r': 1., 'z': 1}

CLASS_COLOR = {'SNIa-norm': 'tab:green', 'SNIbc': 'tab:orange', 'SNII': 'tab:blue', 'SNIIn': 'blue',
               'SNIa-91bg': 'tab:red', 'SNIa-x': 'bisque', 'point-Ia': 'tab:brown', 'Kilonova': '#aaffc3',
               'SLSN-I': 'tab:olive', 'PISN': 'tab:cyan', 'ILOT': '#FF1493', 'CART': 'navy', 'TDE': 'tab:pink',
               'AGN': 'tab:purple'}

# npred = 7
# model_filepath_onepoint_infuture = "/Users/danmuth/OneDrive - University of Cambridge/PycharmProjects/transomaly/plots/model__ci()_ns1_c(1,)/keras_model_epochs300_onepoint_pred7timesteps_infuture_normalised_predict_last49_timesteps_nodropout_100lstmneurons/keras_model_epochs300_onepoint_pred7timesteps_infuture_normalised_predict_last49_timesteps_nodropout_100lstmneurons.hdf5"
model_filepath = "/Users/danmuth/OneDrive - University of Cambridge/PycharmProjects/transomaly/plots/model__ci()_ns1_c(1,)/keras_model_epochs500_pred2timesteps_normalised_predict_last49_timesteps_nodropout_100lstmneurons/keras_model_epochs500_pred2timesteps_normalised_predict_last49_timesteps_nodropout_100lstmneurons.hdf5"

model = load_model(model_filepath, custom_objects={'loss': mean_squared_error()})

passbands = ('g','r')
contextual_info = ()

X = np.array([np.array([[0.09158034, 0.07176773],
       [0.09008677, 0.07137485],
       [0.08917016, 0.0727186 ],
       [0.09029362, 0.07340094],
       [0.08999084, 0.07224263],
       [0.08900606, 0.07197019],
       [0.08739904, 0.07095805],
       [0.08682939, 0.07071227],
       [0.08567506, 0.07031706],
       [0.08386102, 0.07021409],
       [0.0836062 , 0.0698832 ],
       [0.08332578, 0.07050169],
       [0.08397429, 0.06937913],
       [0.0860095 , 0.0716971 ],
       [0.09076107, 0.07780995],
       [0.09871331, 0.08934081],
       [0.11663096, 0.11012718],
       [0.15700709, 0.15327507],
       [0.23417453, 0.22872161],
       [0.35862168, 0.33963544],
       [0.51792838, 0.48456903],
       [0.68616083, 0.64401235],
       [0.81866526, 0.79203408],
       [0.90768435, 0.89154676],
       [0.93726188, 0.94250969],
       [0.92312219, 0.95238174],
       [0.85896026, 0.93542825],
       [0.74499556, 0.87812468],
       [0.62114839, 0.78786493],
       [0.50529031, 0.6932507 ],
       [0.40645818, 0.6036984 ],
       [0.32993884, 0.52479434],
       [0.27163636, 0.45807863],
       [0.22780797, 0.40140124],
       [0.19817062, 0.35339631],
       [0.17674938, 0.31370617],
       [0.16142444, 0.27973599],
       [0.14923072, 0.25118097],
       [0.13910386, 0.22674725],
       [0.13288632, 0.20717057],
       [0.12727512, 0.191151  ],
       [0.12253798, 0.17674126],
       [0.11748478, 0.16368262],
       [0.11428479, 0.15247652],
       [0.110435  , 0.14285344],
       [0.10726069, 0.13410306],
       [0.10317767, 0.12546451],
       [0.10019489, 0.11729226],
       [0.09654251, 0.11025106]])])
y = np.array([np.array([[0.09008677, 0.07137485],
       [0.08917016, 0.0727186 ],
       [0.09029362, 0.07340094],
       [0.08999084, 0.07224263],
       [0.08900606, 0.07197019],
       [0.08739904, 0.07095805],
       [0.08682939, 0.07071227],
       [0.08567506, 0.07031706],
       [0.08386102, 0.07021409],
       [0.0836062 , 0.0698832 ],
       [0.08332578, 0.07050169],
       [0.08397429, 0.06937913],
       [0.0860095 , 0.0716971 ],
       [0.09076107, 0.07780995],
       [0.09871331, 0.08934081],
       [0.11663096, 0.11012718],
       [0.15700709, 0.15327507],
       [0.23417453, 0.22872161],
       [0.35862168, 0.33963544],
       [0.51792838, 0.48456903],
       [0.68616083, 0.64401235],
       [0.81866526, 0.79203408],
       [0.90768435, 0.89154676],
       [0.93726188, 0.94250969],
       [0.92312219, 0.95238174],
       [0.85896026, 0.93542825],
       [0.74499556, 0.87812468],
       [0.62114839, 0.78786493],
       [0.50529031, 0.6932507 ],
       [0.40645818, 0.6036984 ],
       [0.32993884, 0.52479434],
       [0.27163636, 0.45807863],
       [0.22780797, 0.40140124],
       [0.19817062, 0.35339631],
       [0.17674938, 0.31370617],
       [0.16142444, 0.27973599],
       [0.14923072, 0.25118097],
       [0.13910386, 0.22674725],
       [0.13288632, 0.20717057],
       [0.12727512, 0.191151  ],
       [0.12253798, 0.17674126],
       [0.11748478, 0.16368262],
       [0.11428479, 0.15247652],
       [0.110435  , 0.14285344],
       [0.10726069, 0.13410306],
       [0.10317767, 0.12546451],
       [0.10019489, 0.11729226],
       [0.09654251, 0.11025106],
       [0.09630747, 0.10437187]])])
timesX = np.array([np.array([-70., -67., -64., -61., -58., -55., -52., -49., -46., -43., -40.,
       -37., -34., -31., -28., -25., -22., -19., -16., -13., -10.,  -7.,
        -4.,  -1.,   2.,   5.,   8.,  11.,  14.,  17.,  20.,  23.,  26.,
        29.,  32.,  35.,  38.,  41.,  44.,  47.,  50.,  53.,  56.,  59.,
        62.,  65.,  68.,  71.,  74.,  77.])])
objids = np.array(['median_Ia'])

X = np.array([np.array([[0.00448964, 0.        ],
       [0.        , 0.02469979],
       [0.00713832, 0.0510595 ],
       [0.02991802, 0.08449011],
       [0.05011106, 0.09195332],
       [0.05702283, 0.07957474],
       [0.05513788, 0.05611958],
       [0.04902122, 0.0294538 ],
       [0.04288751, 0.00993177],
       [0.04127183, 0.01818543],
       [0.05353569, 0.04212558],
       [0.06324364, 0.06195425],
       [0.05741975, 0.07571231],
       [0.03920042, 0.08415738],
       [0.02232076, 0.08873083],
       [0.01912765, 0.08509039],
       [0.03159054, 0.07017496],
       [0.04909139, 0.05580469],
       [0.06360378, 0.05487265],
       [0.08042658, 0.08175915],
       [0.14612581, 0.15545786],
       [0.32463802, 0.29112003],
       [0.5616513 , 0.4847463 ],
       [0.78721409, 0.70134517],
       [0.94289831, 0.88820695],
       [1.        , 0.98685417],
       [0.9617599 , 1.        ],
       [0.84395464, 0.95538324],
       [0.70269411, 0.87417796],
       [0.56370832, 0.77777756],
       [0.43422203, 0.6788543 ],
       [0.31954263, 0.58431753],
       [0.22572829, 0.5003608 ],
       [0.15997532, 0.43230079],
       [0.13003542, 0.3766077 ],
       [0.11756719, 0.32962512],
       [0.1001643 , 0.29821954],
       [0.08371692, 0.28562479],
       [0.07356015, 0.27645137],
       [0.06571214, 0.2572178 ],
       [0.05947788, 0.23700196],
       [0.05453582, 0.21703187],
       [0.05062684, 0.19788037],
       [0.04754242, 0.17990707],
       [0.04511494, 0.16331596],
       [0.04320991, 0.14819947],
       [0.04171951, 0.1345717 ],
       [0.04055748, 0.12239351],
       [0.03965491, 0.11159111]])])
y = np.array([np.array([[0.        , 0.02469979],
       [0.00713832, 0.0510595 ],
       [0.02991802, 0.08449011],
       [0.05011106, 0.09195332],
       [0.05702283, 0.07957474],
       [0.05513788, 0.05611958],
       [0.04902122, 0.0294538 ],
       [0.04288751, 0.00993177],
       [0.04127183, 0.01818543],
       [0.05353569, 0.04212558],
       [0.06324364, 0.06195425],
       [0.05741975, 0.07571231],
       [0.03920042, 0.08415738],
       [0.02232076, 0.08873083],
       [0.01912765, 0.08509039],
       [0.03159054, 0.07017496],
       [0.04909139, 0.05580469],
       [0.06360378, 0.05487265],
       [0.08042658, 0.08175915],
       [0.14612581, 0.15545786],
       [0.32463802, 0.29112003],
       [0.5616513 , 0.4847463 ],
       [0.78721409, 0.70134517],
       [0.94289831, 0.88820695],
       [1.        , 0.98685417],
       [0.9617599 , 1.        ],
       [0.84395464, 0.95538324],
       [0.70269411, 0.87417796],
       [0.56370832, 0.77777756],
       [0.43422203, 0.6788543 ],
       [0.31954263, 0.58431753],
       [0.22572829, 0.5003608 ],
       [0.15997532, 0.43230079],
       [0.13003542, 0.3766077 ],
       [0.11756719, 0.32962512],
       [0.1001643 , 0.29821954],
       [0.08371692, 0.28562479],
       [0.07356015, 0.27645137],
       [0.06571214, 0.2572178 ],
       [0.05947788, 0.23700196],
       [0.05453582, 0.21703187],
       [0.05062684, 0.19788037],
       [0.04754242, 0.17990707],
       [0.04511494, 0.16331596],
       [0.04320991, 0.14819947],
       [0.04171951, 0.1345717 ],
       [0.04055748, 0.12239351],
       [0.03965491, 0.11159111],
       [0.03895684, 0.10206975]])])
timesX = np.array([np.array([-70., -67., -64., -61., -58., -55., -52., -49., -46., -43., -40.,
       -37., -34., -31., -28., -25., -22., -19., -16., -13., -10.,  -7.,
        -4.,  -1.,   2.,   5.,   8.,  11.,  14.,  17.,  20.,  23.,  26.,
        29.,  32.,  35.,  38.,  41.,  44.,  47.,  50.,  53.,  56.,  59.,
        62.,  65.,  68.,  71.,  74.,  77.])])
objids = np.array(['1_6222236'])

# X[0,:,1][25:] = 0
# X[0,:,0][25:] = 0
# y[0,:,1][24:] = 0
# y[0,:,0][24:] = 0
# X[0,:,2] = 0.04

npred = 2
tidx = -npred+1 if npred != 1 else None

y_pred = model.predict(X[:,:tidx,:2])

idx = 0
sidx = 0


for p in range(50-npred):
       # y_pred_correct_format = np.zeros(y.shape)
       # y_pred_correct_format[:, :, 0] = np.copy(y_predict[:, p, :49])
       # y_pred_correct_format[:, :, 1] = np.copy(y_predict[:, p, :])

       fig, (ax1) = plt.subplots(nrows=1, ncols=1)

       for pbidx, pb in enumerate(passbands):
           plotlabeltest = "ytest:{}".format(pb)
           plotlabelpred = "ypred:{}".format(pb)
           marker = None  # MARKPB[pb] if s == 0 else None
           ax1.plot(timesX[sidx][1:], y[sidx][:, pbidx], c=COLPB[pb], lw=1,
                    label=plotlabeltest, marker=None, markersize=10, alpha=1, linestyle='-')
           ax1.plot(timesX[sidx][1:tidx], y_pred[sidx][:, pbidx+2], c=COLPB[pb], lw=1,
                    label=plotlabelpred, marker=None, markersize=10, alpha=1, linestyle=':')
           # ax1.plot(timesX[sidx][(1+p):(1+p+npred)], y_pred[sidx][p][pbidx::2], c=COLPB[pb], lw=1,
           #          label=plotlabelpred, marker=None, markersize=10, alpha=1, linestyle=':')
           ax1.axvspan(timesX[sidx][1+p], timesX[sidx][-1], alpha=0.2, color='grey')

           # pbmask = lc['passband'] == pb
           # sortedidx = np.argsort(lc[pbmask]['time'].data)
           # time = lc[pbmask]['time'].data[sortedidx]
           # flux = lc[pbmask]['flux'].data[sortedidx]
           # fluxerr = lc[pbmask]['fluxErr'].data[sortedidx]

           # ax1.errorbar(time, flux, yerr=fluxerr,
           #              fmt=".", capsize=0, color=COLPB[pb], label='_nolegend_')

       ax1.set_ylabel("Relative flux")
       ax1.set_xlabel("Time since trigger [days]")
       ax1.legend()
       figdir = 'interesting plots examining what RNN has learned/predict_npred_timesteps_ateachstep/'
       plt.savefig(figdir+f'predict-{npred}-timesteps-{objids[0]}_ateachstep_{p}.png')
       # plt.show()







# npred = 7 # model_filepath_onepoint_infuture = "/Users/danmuth/OneDrive - University of Cambridge/PycharmProjects/transomaly/plots/model__ci()_ns1_c(1,)/keras_model_epochs300_onepoint_pred7timesteps_infuture_normalised_predict_last49_timesteps_nodropout_100lstmneurons/keras_model_epochs300_onepoint_pred7timesteps_infuture_normalised_predict_last49_timesteps_nodropout_100lstmneurons.hdf5" # X[0,:,1][25:] = 0 # X[0,:,0][25:] = 0 # y[0,:,1][24:] = 0 # y[0,:,0][24:] = 0 # X[0,:,2] = 0.04 # y_pred_correct_format = np.zeros(y.shape) # y_pred_correct_format[:, :, 0] = np.copy(y_predict[:, p, :49]) # y_pred_correct_format[:, :, 1] = np.copy(y_predict[:, p, :]) # MARKPB[pb] if s == 0 else None # ax1.plot(timesX[sidx][(1+p):(1+p+npred)], y_pred[sidx][p][pbidx::2], c=COLPB[pb], lw=1, #          label=plotlabelpred, marker=None, markersize=10, alpha=1, linestyle=':') # pbmask = lc['passband'] == pb # sortedidx = np.argsort(lc[pbmask]['time'].data) # time = lc[pbmask]['time'].data[sortedidx] # flux = lc[pbmask]['flux'].data[sortedidx] # fluxerr = lc[pbmask]['fluxErr'].data[sortedidx] # ax1.errorbar(time, flux, yerr=fluxerr, #              fmt=".", capsize=0, color=COLPB[pb], label='_nolegend_') # plt.show()