# coding: utf-8

import os
import glob
from functools import partial
from pp_params import *
import f90nml


class RunSelector:
    def __init__(
        self,
        path_in,
        in_runs=None,
        in_nums="all",
        pp_params=default_params(),
        name_run="*",
        namelist_cond={},
        sort_run_by=None,
        time_min=None,
        time_max=None,
    ):
        self.path_in = path_in
        self.pp_params = pp_params

        self.namelist = {}
        self.runs = self.get_runs(in_runs, name_run, namelist_cond, sort_run_by)

        self.info = {}
        for run in self.runs:
            self.info[run] = {}
        self.nums = {}

        if not type(in_nums) == dict:
            nums_temp = in_nums
            in_nums = {}
            for run in self.runs:
                in_nums[run] = nums_temp

        for run in self.runs:
            self.nums[run] = self.get_nums(run, in_nums[run], time_min, time_max)

    def load_namelist(self, run):
        path_run = self.path_in + "/" + run
        path_nml = path_run + "/" + self.pp_params.input.nml_filename
        return f90nml.read(path_nml)

    def get_nml_value(self, nml_key, run):
        res = self.namelist[run]
        for key in nml_key.split("/"):
            if key in res:
                res = res[key]
            elif key == nml_key.split("/")[-1]:
                res = None
            else:
                raise KeyError(key)
        return res

    def get_runs(self, in_runs=None, name_run="*", namelist_cond={}, sort_run_by=None):
        def try_load_nml(run):
            try:
                self.namelist[run] = self.load_namelist(run)
                success = True
            except IOError:
                success = False
            return success

        runs = map(
            os.path.basename,
            filter(os.path.isdir, glob.glob(self.path_in + "/" + name_run)),
        )
        if not in_runs is None:
            runs = filter(lambda n: n in runs, in_runs)
        runs = filter(try_load_nml, runs)

        if type(namelist_cond) == tuple:
            namelist_cond = [namelist_cond]

        for (nml_key, operator, operand) in namelist_cond:
            value = {}
            for run in runs:
                value[run] = self.get_nml_value(nml_key, run)
            if operator == "=":
                runs = filter(lambda r: value[r] == operand, runs)
            if operator == "!=":
                runs = filter(lambda r: not value[r] == operand, runs)
            elif operator == ">":
                runs = filter(lambda r: value[r] > operand, runs)
            elif operator == "<":
                runs = filter(lambda r: value[r] < operand, runs)
            elif operator == "in":
                runs = filter(lambda r: value[r] in operand, runs)

        # Sort by the value in the namelist of sort_run_by
        if not sort_run_by is None:
            if type(sort_run_by) == str:
                sort_run_by = [sort_run_by]
            for nml_key in reversed(sort_run_by):
                runs.sort(key=partial(self.get_nml_value, nml_key))

        return runs

    def load_info(self, run, num):
        info_file = open(
            self.path_in
            + "/"
            + run
            + "/"
            + "output_"
            + str(num).zfill(5)
            + "/"
            + "info_"
            + str(num).zfill(5)
            + ".txt",
            "r",
        )
        info = {}
        for line in info_file.readlines():
            parsed = yaml.safe_load(line.replace("=", ":"))
            if type(parsed) == dict:
                info.update(parsed)
        info_file.close()
        return info

    def get_nums(self, run, in_nums=None, time_min=None, time_max=None):
        def try_load_info(num):
            try:
                self.info[run][num] = self.load_info(run, num)
                success = True
            except IOError:
                success = False
            return success

        names = glob.glob(
            self.path_in + "/" + run + "/output_[0-9][0-9][0-9][0-9][0-9]"
        )
        nums = map(lambda n: int(n.split("/")[-1].split("_")[1]), names)

        if type(in_nums) == int:
            in_nums = [in_nums]
        if type(in_nums) == list:
            nums = filter(lambda n: n in nums, in_nums)

        nums = np.sort(nums)

        if in_nums == "first":
            i = 0
            while i < len(nums) - 1 and not try_load_info(nums[i]):
                i = i + 1
            nums = [nums[i]]
        elif in_nums == "last":
            i = len(nums) - 1
            while i > 0 and not try_load_info(nums[i]):
                i = i - 1
            nums = [nums[i]]
        else:
            nums = filter(try_load_info, nums)

        if not time_min is None:
            nums = filter(lambda n: self.info[run][n]["time"] >= time_min, nums)
        if not time_max is None:
            nums = filter(lambda n: self.info[run][n]["time"] <= time_max, nums)

        return nums