import fire
import numpy as np 

from utils import *
#from nn import compile_b
from constants import INPUTS
from model import Model
from draw import draw

def predicted_map(B, change=0, path=None):
    year = MAX_YEAR - 1

    df = pd.read_pickle('data.p')

    logger.info('temperature change of %s', change)

    inputs = list(INPUTS)

    for season in SEASONS:
        inputs += [
            'temp_{}_{}'.format(season, year),
            'precip_{}_{}'.format(season, year)
        ]

    frame = df[inputs + ['longitude']]
    frame_cp = df[inputs + ['longitude']]

    for season in SEASONS:
        frame.loc[:, 'temp_{}_{}'.format(season, year)] += change

    columns = ['latitude', 'longitude', 'biome_num']
    new_data = pd.DataFrame(columns=columns)
    nframe = pd.DataFrame(columns=frame.columns, data=normalize_ndarray(frame.to_numpy(), frame_cp.to_numpy()))

    for i, (chunk, chunk_original) in enumerate(zip(chunker(nframe, B.batch_size), chunker(frame, B.batch_size))):
        if chunk.shape[0] < B.batch_size:
            continue
        input_data = chunk.loc[:, inputs].values
        out = B.predict(input_data)

        f = pd.DataFrame({
            'longitude': chunk_original.loc[:, 'longitude'], 
            'latitude': chunk_original.loc[:, 'latitude'],
            'biome_num': out
        }, columns=columns)
        new_data = new_data.append(f)

    draw(new_data, path=path)

def predicted_map_cmd(checkpoint='checkpoints/save.h5', change=0, path=None):
    B = Model('b', epochs=1)
    B.prepare_for_use()
    B.restore(checkpoint)
    predicted_map(B, change=change, path=path)

if __name__ == "__main__":
    fire.Fire(predicted_map_cmd)