import numpy as np 

from utils import *
from nn import B
from draw import draw
import time

def chunker(seq, size):
    return (seq[pos:pos + size] for pos in range(0, len(seq), size))

year = MAX_YEAR - 1

df = pd.read_pickle('data_final.p')
latitude = np.array(df.index.get_level_values(1))
df.loc[:, 'latitude'] = pd.Series(latitude, index=df.index)

B.prepare_dataset(df, dataframe_to_dataset_biomes)
B.create_model([64, 128])
B.compile()

for change in range(-5, 6):
    print('TEMPERATURE MODIFICATION OF {}'.format(change))

    inputs = ['elevation', 'distance_to_water']

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

    inputs += ['latitude']

    frame = df[inputs]
    print(frame.head())

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

    # print(np.average(frame.loc[:, 'temp_winter_2016']))

    # index = []
    # for longitude in range(-179, 179):
        # for latitude in range(-89, 89):
            # index.append((longitude, latitude))

    columns = ['biome_num']
    new_data = pd.DataFrame(columns=columns)
    for i, chunk in enumerate(chunker(frame, B.batch_size)):
        input_data = normalize_ndarray(chunk.values)
        out = B.predict(input_data)
        new_index = np.concatenate((chunk.index.values, new_data.index.values))

        new_data = new_data.reindex(new_index)
        new_data.loc[chunk.index.values, 'biome_num'] = out

    # print(new_data['biome_num'].unique())

    draw(new_data)

    # columns = ['biome_num']
    # indices = ['longitude', 'latitude']
    # new_df = pd.DataFrame(index=indices, columns=columns)
    # new_df =