# coding: utf-8

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


     def __init__(self, path_in,
                  in_runs=None, in_nums="all", pp_params=default_params(),
                  name_run='*', namelist_cond = dict(),
                  sort_run_by=None, time_min=None, time_max=None):
          self.path_in = path_in
          self.pp_params = pp_params

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

          self.info = dict()
          for run in self.runs:
               self.info[run] = dict()
          self.nums = dict()

          if not type(in_nums) == dict :
               nums_temp = in_nums
               in_nums = dict()
               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=dict(), 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 = dict()
               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 = dict()
          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