import numpy as np
import imageio
import pickle
import tensorflow as tf
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import matplotlib.ticker as mticker
from cartopy.mpl.gridliner import LATITUDE_FORMATTER, LONGITUDE_FORMATTER
import shapefile as shp
from matplotlib import animation
from IPython.display import HTML

def get_array(source, scaler_dict):

    reader = imageio.get_reader(source)

    source_video = []

    try:
        for im in reader:
            source_video.append(im)
    except RuntimeError:
        pass

    reader.close()

    scaler_path = scaler_dict[source[-14:]]

    with open(scaler_path, 'rb') as f:
        sc = pickle.load(f)

    data = np.array(source_video)[:,:,:,0]
    data = sc.inverse_transform(data)
    data = np.swapaxes(data, 0, 2)
    data = np.swapaxes(data, 0, 1)

    X = data[:,:,0:12]
    y = data[:,:,12:]

    return X, y

def get_slices(values, slices):

    dim_size = len(values)
    idx_step = int(dim_size/slices)

    slices_list = []

    for i in range(idx_step, dim_size, idx_step):

        slices_list.append(np.round(values[i], 2))
    
    return slices_list

def save_video(X, threshold=0, file_path = 'data/video.mp4'):

    # Get vmax
    var = X.copy()

    var[np.isnan(var)] = 0
    var[var<=0] = 0

    counts, bins = np.histogram(var[:])

    value = counts[counts>np.median(counts)][-1]
    idx = np.where(counts==value)[0][0]

    vmax = np.round(bins[idx])
    
    # Latitude and longitude
    lon   = np.loadtxt('data/longitude.txt')
    lat   = np.loadtxt('data/latitude.txt')

    area  = [lon.min(),lon.max(),lat.min(),lat.max()]

    lat_list = get_slices(lat, 4)
    lon_list = get_slices(lon, 6)

    # Visualization
    ims = []
    fig = plt.figure(figsize=(7,5))
    ax = plt.axes(projection=ccrs.PlateCarree())

    gl = ax.gridlines(crs=ccrs.PlateCarree(),
                    draw_labels=True,
                    linewidth=0.3,
                    color='black',
                    linestyle='--')

    gl.top_labels = False
    gl.right_labels = False
    gl.xlines = True
    gl.xlocator = mticker.FixedLocator(lon_list)
    gl.ylocator = mticker.FixedLocator(lat_list)
    gl.xformatter = LONGITUDE_FORMATTER
    gl.yformatter = LATITUDE_FORMATTER
    gl.xlabel_style = {'size':10, 'color':'black'}
    gl.ylabel_style = {'size':10, 'color':'black'}

    frames = X
    frames[frames<=threshold] = np.nan

    barra = np.arange(0, vmax+1, 5)

    for i in range(frames.shape[2]):
        im = plt.imshow(frames[..., i],
                        cmap=plt.cm.rainbow,
                        vmin=0,
                        vmax=vmax,
                        extent=area,
                        origin='lower',
                        animated=True)
        
        ims.append([im])

    cbar = plt.colorbar(ax=ax, pad=0.02, aspect=16, shrink=0.77)
    cbar.set_ticks(barra)
    cbar.set_label('mm/h')

    shapeID = shp.Reader("data/shapefile/regiao_sul.shp")

    for shape in shapeID.shapeRecords():
        point = np.array( shape.shape.points )
        dummy = plt.plot( point[:,0] , point[:,1], color="black", linewidth=0.5 ) # 1

    ani = animation.ArtistAnimation(fig, ims, interval=500, blit=True, repeat_delay=1000)

    FFwriter = animation.FFMpegWriter(fps=2)
    ani.save(file_path, writer = FFwriter)

    # ani.save(f'data/vis.gif', writer='pillow', fps=6)

    plt.close(ani._fig)
    HTML(ani.to_html5_video())

def make_predictions(X):

    filepath = "models/model.h5"
    model = tf.keras.models.load_model(filepath)

    X[np.isnan(X)] = 0

    X = np.expand_dims(X, axis=0)

    scaler_path = 'models/scaler.pkl'
    with open(scaler_path, 'rb') as f:
        sc = pickle.load(f)

    X = sc.transform(X)
    ypred = model.predict(X)
    print(ypred.shape)
    ypred = sc.inverse_transform(ypred)[0]
    print(ypred.shape)

    return ypred