import random
import numpy as np
import itertools
from functools import reduce
from dataclasses import asdict
from typing import Tuple
from .ParameterDefinitions import Parameters, JsonFormat
from .GraphicUserInterface import Visualizer
from .Network import Server, request_abaqus
from .FileIO import pickle_io, pickles_io, remove_file, get_sorted_file_numbers_from_pattern
from .PostProcessing import evaluate_all_fitness_values, selection
from .MutateAndValidate import mutate_and_validate_topology


class NSGAModel:
    def __init__(self, params: Parameters, material_properties: dict, fitness_definitions: dict,
                 visualizer: Visualizer = None, random_topology_density: float = 0.5):
        self.params = params
        self.fitness_definitions = fitness_definitions
        self.visualizer = visualizer
        self.material_properties = material_properties
        self.random_topology_density = random_topology_density

    def load_parent_data(self, gen: int, server: Server) -> Tuple[np.ndarray, dict]:
        try:
            parent_topologies = pickle_io(f'Topologies_{gen}', mode='r')['parent']
        except FileNotFoundError or KeyError:
            parent_topologies = random_parent_generation(density=self.random_topology_density,
                                                         params=self.params, save_file_as=f'Topologies_{gen}')
        try:
            parent_results = pickle_io(f'FieldOutput_{gen}', mode='r')
        except FileNotFoundError:
            json_data = asdict(JsonFormat(start_topology_from=1, topologies_key='parent',
                                          topologies_file_name=f'Topologies_{gen}', exit_abaqus=False))
            json_data.update(asdict(self.params))
            json_data.update(self.material_properties)
            request_abaqus(dict_data=json_data, server=server, conn_to_gui=self.visualizer.conn_to_gui)
            parent_results = pickle_io(f'FieldOutput_offspring_{gen}', mode='r')
            remove_file(f'FieldOutput_offspring_{gen}')
            pickle_io(f'FieldOutput_{gen}', mode='w', to_dump=parent_results)
            self.visualizer.visualize(params=self.params, gen=gen - 1, use_manual_rp=False)
        assert len(parent_topologies) == len(parent_results)
        return parent_topologies, parent_results

    def determine_where_abaqus_start(self) -> Tuple[int, int]:
        offspring_results_file_numbers = get_sorted_file_numbers_from_pattern(r'FieldOutput_offspring_\d+')
        if len(offspring_results_file_numbers) == 0:
            return 1, 1
        else:
            last_offspring_results_file_number = offspring_results_file_numbers[-1]
            last_offspring_results = pickle_io(f'FieldOutput_offspring_{last_offspring_results_file_number}', mode='r')
            if len(last_offspring_results) < self.params.end_pop:
                needed_to_start_at_offspring = len(last_offspring_results) + 1
                return last_offspring_results_file_number, needed_to_start_at_offspring
            else:
                return last_offspring_results_file_number + 1, 1

    def generate_offspring_topologies(self, gen: int, server: Server) -> np.ndarray:
        parent_topologies, parent_results = self.load_parent_data(gen=gen, server=server)
        topologies = pickle_io(f'Topologies_{gen}', mode='r')
        if 'offspring' in topologies.keys():
            offspring_topologies = pickle_io(f'Topologies_{gen}', mode='r')['offspring']
        else:
            offspring_topologies = generate_offspring(gen=gen, topo_parents=parent_topologies, params=self.params,
                                                      save_file_as=f'Topologies_{gen}')
        assert len(parent_topologies) == len(parent_results) == len(offspring_topologies)
        return offspring_topologies

    def evolve_a_generation(self, running_gen: int, start_offspring_from: int, server: Server):  # changed method name: from .run_a_generation() to .evolve_a_generation()
        parent_topologies, parent_results = self.load_parent_data(gen=running_gen, server=server)
        offspring_topologies = self.generate_offspring_topologies(gen=running_gen, server=server)
        json_data = asdict(JsonFormat(start_topology_from=start_offspring_from, topologies_key='offspring',
                                      topologies_file_name=f'Topologies_{running_gen}', exit_abaqus=False))
        json_data.update(asdict(self.params))
        json_data.update(self.material_properties)
        request_abaqus(dict_data=json_data, server=server, conn_to_gui=self.visualizer.conn_to_gui)
        offspring_results = pickle_io(f'FieldOutput_offspring_{running_gen}', mode='r')
        all_topologies = np.vstack((parent_topologies, offspring_topologies))
        all_results = parent_results.copy()
        all_results.update({entity_num + len(parent_results): offspring_results[entity_num]
                            for entity_num in sorted(offspring_results.keys())})
        all_fitness_values = evaluate_all_fitness_values(fitness_definitions=self.fitness_definitions,
                                                         params_dict=asdict(self.params),
                                                         results=all_results, topologies=all_topologies)
        pareto_indices = selection(all_fitness_values=all_fitness_values, selected_size=self.params.end_pop)
        selected_topologies = all_topologies[pareto_indices]
        selected_results = {entity_num: all_results[pareto_idx + 1]
                            for entity_num, pareto_idx in enumerate(pareto_indices, start=1)}
        assert len(selected_topologies) == len(selected_results)
        pickle_io(f'Topologies_{running_gen + 1}', mode='w', to_dump={'parent': selected_topologies})
        pickle_io(f'FieldOutput_{running_gen + 1}', mode='w', to_dump=selected_results)
        if self.visualizer is not None:
            self.visualizer.visualize(params=self.params, gen=running_gen, use_manual_rp=False)

    def evolve(self, server):  # changed method name: from .run() to .evolve()
        start_gen, start_offspring = self.determine_where_abaqus_start()
        for gen in range(start_gen, self.params.end_gen):
            self.evolve_a_generation(running_gen=gen, start_offspring_from=start_offspring, server=server)
            start_offspring = 1
        request_abaqus(dict_data={'exit_abaqus': True}, server=server, conn_to_gui=self.visualizer.conn_to_gui)


def find_where_same_array_locates(arr_to_find: np.ndarray, big_arr: np.ndarray) -> np.ndarray:
    axis_range = tuple(dim for dim in range(len(big_arr.shape) - len(arr_to_find.shape), len(big_arr.shape)))
    return np.argwhere(np.all(big_arr == arr_to_find, axis=axis_range))


def get_cutting_section_and_candidates(topologies: np.ndarray) -> Tuple[int, np.ndarray]:
    topologies_flattened = topologies.reshape((len(topologies), -1))
    possible_cutting_sections = np.random.permutation(topologies_flattened.shape[1])
    for possible_cutting_section in possible_cutting_sections:
        candidates_for_cutting = np.empty((0,), dtype=int)
        for topo_idx in range(len(topologies)):
            if topologies_flattened[topo_idx, possible_cutting_section] == 1:
                candidates_for_cutting = np.hstack((candidates_for_cutting, topo_idx))
        if len(candidates_for_cutting) >= 2:
            return possible_cutting_section, candidates_for_cutting
    raise RuntimeError


def get_candidate_pairs(candidates):
    candidates_results = list(itertools.combinations(candidates, 2))  # possible candidate pair of parents
    random.shuffle(candidates_results)
    return candidates_results


def crossover(chromosome_1, chromosome_2, cutting_section):  # Crossover process
    chromosome_1_flattened = chromosome_1.flatten()
    chromosome_2_flattened = chromosome_2.flatten()
    cross_overed_chromosome_1 = np.empty_like(chromosome_1_flattened)
    cross_overed_chromosome_2 = np.empty_like(chromosome_2_flattened)
    cross_overed_chromosome_1[0:cutting_section] = chromosome_1_flattened[0:cutting_section]
    cross_overed_chromosome_1[cutting_section:] = chromosome_2_flattened[cutting_section:]
    cross_overed_chromosome_2[0:cutting_section] = chromosome_2_flattened[0:cutting_section]
    cross_overed_chromosome_2[cutting_section:] = chromosome_1_flattened[cutting_section:]
    return cross_overed_chromosome_1.reshape(chromosome_1.shape), cross_overed_chromosome_2.reshape(chromosome_2.shape)


def generate_offspring(gen: int, params: Parameters, topo_parents: np.ndarray, save_file_as: str = None) -> np.ndarray:
    """
    Generating topo_offspring from parents. Crossover & Mutation & Validating processes will be held.
    Validating processes contain, checking 3d-print-ability without support, one voxel tree contacting six faces
    in a cube
    :param save_file_as: Save pickle file of offspring topologies as this.
    :param topo_parents: Topology array of parent, shape: (end_pop, lx, ly, lz)
    :param gen: Current generation
    :param params: Parameters for GA
    The more timeout, the longer time will be allowed for the function "mutate_and_validate_topology".
    :return:
    """
    topo_offspring = np.empty((0, params.lx, params.ly, params.lz), int)
    validation_count = 0
    all_topos = pickles_io(file_names=[f'Topologies_{g}' for g in range(1, gen + 1)], mode='r', key_option='int')
    all_topos_parent = np.array([topos['parent'] for topos in all_topos.values()], dtype=int)
    while True:
        cutting_section, candidates = get_cutting_section_and_candidates(topologies=topo_parents)
        candidate_pairs = get_candidate_pairs(candidates=candidates)
        print('<info> Candidate lists:', len(candidates))
        print('<info> Candidate pairs:', len(candidate_pairs))
        for pair_idx in range(len(candidate_pairs)):
            chromosome_1_idx = candidate_pairs[pair_idx][0]
            chromosome_2_idx = candidate_pairs[pair_idx][1]
            cross_overed_chromosome_1, cross_overed_chromosome_2 = crossover(
                chromosome_1=topo_parents[chromosome_1_idx],
                chromosome_2=topo_parents[chromosome_2_idx],
                cutting_section=cutting_section)
            validated_chromosome_1 = mutate_and_validate_topology(cross_overed_chromosome_1, params.mutation_rate)
            validated_chromosome_2 = mutate_and_validate_topology(cross_overed_chromosome_2, params.mutation_rate)
            for validated_chromosome in (validated_chromosome_1, validated_chromosome_2):
                if validated_chromosome is None:
                    print('<!> Non-connected tree detected')
                    continue
                else:
                    if len(find_where_same_array_locates(arr_to_find=validated_chromosome,
                                                         big_arr=all_topos_parent)):
                        print('<!> Clone structure found in parents!')
                        continue
                    elif len(find_where_same_array_locates(arr_to_find=validated_chromosome,
                                                           big_arr=topo_offspring)):
                        print('<!> Clone structure found in current offsprings!')
                        continue
                    else:
                        topo_offspring = np.vstack((topo_offspring, np.expand_dims(validated_chromosome, axis=0)))
                        validation_count += 1
                        print(f'<info> Validation of chromosome {validation_count} complete!')
                        if len(topo_offspring) == params.end_pop:
                            print('<info> Generating topo_offspring complete')
                            if save_file_as is not None:
                                pickle_io(save_file_as, mode='a', to_dump={'offspring': topo_offspring})
                            return topo_offspring


def random_array(shape, probability):
    return np.random.choice([1, 0], size=reduce(lambda x, y: x * y, shape), p=[probability, 1 - probability]).reshape(
        shape)


def random_parent_generation(density: float, params: Parameters, save_file_as: str = None) -> np.ndarray:
    parents = np.empty((params.end_pop, params.lx, params.ly, params.lz))
    total_parent_generation_count = 0
    total_volume_frac = 0
    while total_parent_generation_count < params.end_pop:
        print(f'<<<<< Parent {total_parent_generation_count + 1} >>>>>')
        while True:
            rand_arr = random_array(shape=(params.lx, params.ly, params.lz), probability=density)
            parent = mutate_and_validate_topology(rand_arr, mutation_probability=params.mutation_rate)
            if parent is not None:
                break
        volume_frac = np.count_nonzero(parent) / (params.lx * params.ly * params.lz / 100)
        total_volume_frac += volume_frac
        print(f'Volume fraction: {volume_frac:.1f} %\n')
        parents[total_parent_generation_count] = parent
        total_parent_generation_count += 1
    print(f'Average volume fraction: {total_volume_frac / params.end_pop:.1f} %')
    if save_file_as is not None:
        pickle_io(save_file_as, mode='w', to_dump={'parent': parents})
    return parents
