import numpy as np
import pandas as pd
import time
import os
import json
from concurrent.futures import ThreadPoolExecutor
from deap import base, creator, tools, algorithms


def fitness_numpy(individual, price, load, temperature, irradiance, wind_speed, prev_soc, prev_Pg1, prev_Pg2, prev_Pg3):
    individual = np.array(individual)
    price = np.array(price)
    load = np.array(load)
    temperature = np.array(temperature)
    irradiance = np.array(irradiance)
    wind_speed = np.array(wind_speed)

    Ac, Ag1, Ag2, Ag3, Av = individual[:5]
    soc = np.clip(prev_soc + 0.2 * Ac * 0.9, 0.2, 0.8)
    Pg1 = np.clip(prev_Pg1 + 100 * Ag1, 0, 150)
    Pg2 = np.clip(prev_Pg2 + 100 * Ag2, 0, 375)
    Pg3 = np.clip(prev_Pg3 + 200 * Ag3, 0, 500)
    Pso = np.clip((0.2 * irradiance + 0.05 * temperature - 9.25) * (1 + Av), 0, None)
    Pw = np.where((wind_speed >= 3) & (wind_speed < 8), wind_speed ** 3 * 172.2625 / 1000,
                  np.where((wind_speed >= 8) & (wind_speed < 12), 64 * 172.2625 / 125, 0))
    P = Ac + Pg1 + Pg2 + Pg3 + Pso + Pw

    Ee = np.where(P >= load, P - load, 0)
    Es = np.where(P < load, load - P, 0)

    Cb = 0.01 * Ac + 0.1 * soc
    Cg1 = 0.0034 * Pg1 ** 2 + 3 * Pg1 + 30
    Cg2 = 0.001 * Pg2 ** 2 + 10 * Pg2 + 40
    Cg3 = 0.001 * Pg3 ** 2 + 15 * Pg3 + 70
    Cs = 0.01 * Pso
    Cw = 0.01 * Pw
    Rs = 0.5 * price * Ee
    Cp = price * Es
    Pe = np.where(Ee > 100, (Ee - 100) * 50, 0)
    Ps = np.where(Es > 100, (Es - 100) * 50, 0)

    total_cost = np.sum(Cb + Cg1 + Cg2 + Cg3 + Cs + Cw - Rs + Cp + Pe + Ps)
    reward = -total_cost / 1000

    return (reward,)


def check_bounds(func):
    def wrapper(*args, **kwargs):
        offspring = func(*args, **kwargs)
        for individual in offspring:
            for i in range(len(individual)):
                individual[i] = np.clip(individual[i], -1, 1)
        return offspring
    return wrapper


def main():
    data = pd.read_csv('./data.csv')

    price = data['price'].values
    load = data['load'].values
    temperature = data['temperature'].values
    irradiance = data['irradiance'].values
    wind_speed = data['wind_speed'].values

    creator.create("FitnessMax", base.Fitness, weights=(1.0,))
    creator.create("Individual", list, fitness=creator.FitnessMax)

    toolbox = base.Toolbox()
    toolbox.register("attr_float", np.random.uniform, -1, 1)
    toolbox.register("individual", tools.initRepeat, creator.Individual, toolbox.attr_float, n=5)
    toolbox.register("population", tools.initRepeat, list, toolbox.individual)

    toolbox.register("mate", check_bounds(tools.cxBlend), alpha=0.5)
    toolbox.register("mutate", check_bounds(tools.mutGaussian), mu=0, sigma=0.1, indpb=0.2)
    toolbox.register("select", tools.selTournament, tournsize=3)
    toolbox.register("evaluate", fitness_numpy)

    population = toolbox.population(n=500)
    prev_soc, prev_Pg1, prev_Pg2, prev_Pg3 = 0.4, 0.0, 0.0, 0.0
    decision_values = []

    for hour in range(8760):
        start = time.time()
        current_price = price[hour]
        current_load = load[hour]
        current_temperature = temperature[hour]
        current_irradiance = irradiance[hour]
        current_wind_speed = wind_speed[hour]

        for gen in range(500):
            offspring = toolbox.select(population, len(population))
            offspring = list(map(toolbox.clone, offspring))

            for child1, child2 in zip(offspring[::2], offspring[1::2]):
                if np.random.rand() < 0.7:
                    toolbox.mate(child1, child2)
                    del child1.fitness.values
                    del child2.fitness.values
            for mutant in offspring:
                if np.random.rand() < 0.2:
                    toolbox.mutate(mutant)
                    del mutant.fitness.values

            invalid_ind = [ind for ind in offspring if not ind.fitness.valid]

            with ThreadPoolExecutor() as executor:
                futures = []
                batch_size = 500

                for i in range(0, len(invalid_ind), batch_size):
                    batch = invalid_ind[i:i + batch_size]
                    for ind in batch:
                        futures.append(
                            executor.submit(toolbox.evaluate, ind, current_price, current_load, current_temperature,
                                            current_irradiance, current_wind_speed, prev_soc, prev_Pg1, prev_Pg2,
                                            prev_Pg3))

                for future in futures:
                    fitness = future.result()
                    for ind in invalid_ind:
                        if not ind.fitness.valid:
                            ind.fitness.values = fitness

            population[:] = offspring
        end = time.time()
        best_ind = tools.selBest(population, 1)[0]
        print(f'Best individual at hour {hour + 1}: {best_ind}')
        print(f'Fitness: {best_ind.fitness.values}, using {end - start}s')

        decision_values.append({
            'Ac': best_ind[0],
            'Ag1': best_ind[1],
            'Ag2': best_ind[2],
            'Ag3': best_ind[3],
            'Av': best_ind[4]
        })

        prev_soc, prev_Pg1, prev_Pg2, prev_Pg3 = best_ind[0], best_ind[1], best_ind[2], best_ind[3]

    with open('decision_values.json', 'w') as f:
        json.dump(decision_values, f)


if __name__ == "__main__":
    main()