import time

import os
from pywmi import evaluate, smt_to_nested

from incal.incremental_learner import IncrementalObserver

from pywmi.plot import plot_combined


class PlottingObserver(IncrementalObserver):
    def __init__(
        self,
        domain,
        directory,
        name,
        feat_x,
        feat_y,
        condition=None,
        auto_clean=False,
        run_name=None,
    ):
        self.domain = domain

        if not os.path.exists(directory):
            os.makedirs(directory)

        if auto_clean:
            run_number = 0
            run_dir = None
            while run_dir is None or os.path.exists(run_dir):
                date_folders = time.strftime("%Y{s}%m{s}%d{s}".format(s=os.path.sep))
                run_name = run_name + " " if run_name is not None else ""
                run_dir_name = "run {}{}".format(run_name, time.strftime("%Hh %Mm %Ss"))
                run_dir = os.path.join(directory, date_folders, run_dir_name)
                if run_number > 0:
                    run_dir += "_{}".format(run_number)
                run_number += 1
            os.makedirs(run_dir)
            directory = run_dir

        self.directory = directory

        self.name = name
        self.all_active = set()
        self.feat_x = feat_x
        self.feat_y = feat_y
        self.iteration = 0
        self.condition = condition

    def observe_initial(self, data, labels, initial_indices):
        self.all_active = self.all_active.union(initial_indices)
        name = "{}{}{}_{}".format(
            self.directory, os.path.sep, self.name, self.iteration
        )
        plot_combined(
            self.feat_x,
            self.feat_y,
            self.domain,
            None,
            (data, labels),
            None,
            name,
            initial_indices,
            set(),
            self.condition,
        )

    def observe_iteration(
        self, data, labels, formula, new_active_indices, solving_time, selection_time
    ):
        self.iteration += 1
        learned_labels = evaluate(self.domain, formula, data)
        name = "{}{}{}_{}".format(
            self.directory, os.path.sep, self.name, self.iteration
        )
        plot_combined(
            self.feat_x,
            self.feat_y,
            self.domain,
            formula,
            (data, labels),
            learned_labels,
            name,
            self.all_active,
            new_active_indices,
            condition=self.condition,
        )
        self.all_active = self.all_active.union(new_active_indices)
