import os
from time import sleep

import numpy as np
from ase.calculators.singlepoint import SinglePointCalculator
from ase.io import read, write

from agox.observer import Observer
from agox.samplers.metropolis import MetropolisSampler


def safe_write(filename, atoms):
    write(filename=filename + ".lock", images=atoms, format="traj")
    os.rename(filename + ".lock", filename)


class ConcurrentTemperingSampler(MetropolisSampler):
    """
    Parallel tempering sampler for concurrent searches with multiple worker jobs. Written to work with ConcurrentDatabase, no guarantee that it works in any other setting.

    Parameters
    ----------
    temperatures : list
        List of temperatures for each worker.
    database : ConcurrentDatabase
        The database that the sampler will attach to.
    swap_frequency : int
        The frequency at which the workers will swap candidates.
    gets : list
        List of dictionaries with the gets that the sampler will observe.
    sets : list
        List of dictionaries with the sets that the sampler will observe.
    swap_order : int
        Order of the observer.
    """

    name = "ParallelTemperingSampler"

    def __init__(
        self,
        temperatures=[],
        database=None,
        swap_frequency=10,
        gets=[{"get_key": "evaluated_candidates"}, {}],
        sets=[{}, {}],
        swap_order=2,
        verbose: bool = False,
        **kwargs,
    ):
        super().__init__(gets=gets, sets=sets, **kwargs)
        self.swap_frequency = swap_frequency
        self.temperatures = temperatures
        self.temperature = temperatures[database.worker_number]
        self.swap_order = swap_order
        self.swap_func = self.metropolis_hastings_swap
        self.most_recently_accepted = {worker_number: None for worker_number in range(database.total_workers)}
        self.most_recently_accepted_iteration = {worker_number: 0 for worker_number in range(database.total_workers)}
        self.verbose = verbose

        self.add_observer_method(
            self.swap_candidates,
            sets=self.sets[1],
            gets=self.gets[1],
            order=self.swap_order,
            handler_identifier="database",
        )

        self.attach_to_database(database)

    @Observer.observer_method
    def swap_candidates(self, database, state):
        """
        Assumes that the database is synced.
        """

        if self.decide_to_swap(database):
            for candidate in database.candidates:
                worker_number = candidate.get_meta_information("worker_number")
                iteration = candidate.get_meta_information("iteration")
                accepted = candidate.get_meta_information("accepted")

                if accepted:
                    if iteration > self.most_recently_accepted_iteration[worker_number]:
                        self.most_recently_accepted[worker_number] = candidate
                        self.most_recently_accepted_iteration[worker_number] = candidate.get_meta_information(
                            "iteration"
                        )

            if self.verbose:
                energies = [
                    self.most_recently_accepted[c].get_potential_energy()
                    if self.most_recently_accepted[c] is not None
                    else 0
                    for c in self.most_recently_accepted
                ]
                if database.worker_number == 0:
                    text = (
                        "PARALLEL TEMPERING:" + f"{self.temperature:8.3f}, " + ",".join([f"{e:8.3f}" for e in energies])
                    )
                else:
                    text = "PARALLEL TEMPERING (not main worker):"
                self.writer(text)

            self.swap_func(database)

    def metropolis_hastings_swap(self, database):
        worker_number = database.worker_number
        total_workers = database.total_workers
        iteration = self.get_iteration_counter()

        # I abuse the filenames a bit here:
        filename = database.filename[:-3] + "_swap_iteration_{}_worker_{}.traj"

        if worker_number == 0:
            # This one does the calculation and 'broadcasts' to the others over disk.

            # Starting from the bottom:
            for i in range(total_workers - 1):
                C_i = self.most_recently_accepted[i]
                C_j = self.most_recently_accepted[i + 1]

                E_i = C_i.get_potential_energy()
                E_j = C_j.get_potential_energy()

                beta_i = 1 / self.temperatures[i]
                beta_j = 1 / self.temperatures[i + 1]

                P = np.min([1, np.exp((beta_i - beta_j) * (E_i - E_j))])

                r = np.random.rand()

                if r < P:
                    self.writer("Swapped {} with {}".format(i, i + 1))
                    self.most_recently_accepted[i], self.most_recently_accepted[i + 1] = (
                        self.most_recently_accepted[i + 1],
                        self.most_recently_accepted[i],
                    )
                    self.most_recently_accepted_iteration[i], self.most_recently_accepted_iteration[i + 1] = (
                        iteration,
                        iteration,
                    )

                else:
                    self.writer("Did not swap {} for {}".format(i, i + 1))

            # Write the candidates:
            for wn in range(1, total_workers):
                atoms = self.most_recently_accepted[wn]
                atoms.info = atoms.meta_information
                safe_write(filename.format(iteration, wn), atoms)
            self.chosen_candidate = self.most_recently_accepted[worker_number]

        else:
            while not os.path.exists(filename.format(iteration, worker_number)):
                sleep(1)

            chosen_atoms = read(filename.format(iteration, worker_number))

            self.sample[0] = self.convert_to_candidate_object(chosen_atoms, self.sample[0].template)
            self.sample[0].meta_information = chosen_atoms.info
            scp = SinglePointCalculator(
                self.sample[0], energy=chosen_atoms.get_potential_energy(), forces=chosen_atoms.get_forces()
            )
            self.sample[0].set_calculator(scp)

            os.remove(filename.format(iteration, worker_number))

    def convert_to_candidate_object(self, atoms_type_object, template):
        candidate = self.candidate_instanstiator(
            template=template,
            positions=atoms_type_object.positions,
            numbers=atoms_type_object.numbers,
            cell=atoms_type_object.cell,
            info=atoms_type_object.info,
        )
        return candidate

    def decide_to_swap(self, database):
        return (self.get_iteration_counter() % self.swap_frequency == 0) * (database.total_workers > 1)
