# -*- coding: utf-8 -*-
# vim: noai:ts=4:sw=4:expandtab

import fcntl
import glob
import grp
import logging
import os
import pwd
import shlex
import shutil
import stat
import tempfile
import uuid

from . import mounts
from . import uid
from . import util
from .exception import BuildRootLocked, Error, ResultDirNotAccessible, RootError
from .package_manager import package_manager
from .trace_decorator import getLog, traceLog

class Buildroot(object):
    @traceLog()
    def __init__(self, config, uid_manager, state, plugins, bootstrap_buildroot=None, is_bootstrap=False):
        self.config = config
        self.uid_manager = uid_manager
        self.state = state
        self.plugins = plugins
        self.bootstrap_buildroot = bootstrap_buildroot
        self.is_bootstrap = is_bootstrap
        self.shared_root_name = config['root']
        if 'unique-ext' in config:
            config['root'] = "%s-%s" % (config['root'], config['unique-ext'])
        self.root_name = config['root']
        self.mockdir = config['basedir']
        self.basedir = os.path.join(config['basedir'], config['root'])
        if 'rootdir' in config:
            self.rootdir = config['rootdir']
        else:
            self.rootdir = os.path.join(self.basedir, 'root')
        self.resultdir = config['resultdir'] % config
        self.homedir = config['chroothome']
        self.cache_topdir = config['cache_topdir']
        self.cachedir = os.path.join(self.cache_topdir, self.shared_root_name)
        self.builddir = os.path.join(self.homedir, 'build')
        self._lock_file = None
        self.selinux = (not self.config['plugin_conf']['selinux_enable']
                        and util.selinuxEnabled())

        self.chrootuid = config['chrootuid']
        self.chrootuser = 'mockbuild'
        self.chrootgid = config['chrootgid']
        self.chrootgroup = config['chrootgroup']
        self.env = config['environment']
        self.env['HOME'] = self.homedir
        proxy_env = util.get_proxy_environment(config)
        self.env.update(proxy_env)
        os.environ.update(proxy_env)

        self.pkg_manager = package_manager(config, self, plugins, bootstrap_buildroot)
        self.mounts = mounts.Mounts(self)

        self.root_log = getLog("mockbuild")
        self.build_log = getLog("mockbuild.Root.build")
        self.logging_initialized = False
        self.chroot_was_initialized = False
        self.preexisting_deps = []
        self.plugins.init_plugins(self)
        self.tmpdir = None
        self.nosync_path = None

    @traceLog()
    def make_chroot_path(self, *paths):
        new_path = self.rootdir
        for path in paths:
            if path.startswith('/'):
                path = path[1:]
            new_path = os.path.join(new_path, path)
        return new_path

    @traceLog()
    def initialize(self, prebuild=False, do_log=True):
        """
        Initialize the buildroot to a point where it's possible to execute
        commands in chroot. If it was already initialized, just lock the shared
        lock.
        """
        try:
            self._lock_buildroot(exclusive=True)
            self._init(prebuild=prebuild, do_log=do_log)
        except BuildRootLocked:
            pass
        finally:
            self._lock_buildroot(exclusive=False)
        if do_log:
            self._resetLogging()

    @traceLog()
    def chroot_is_initialized(self):
        return os.path.exists(self.make_chroot_path('.initialized'))

    @traceLog()
    def _setup_result_dir(self):
        with self.uid_manager:
            try:
                util.mkdirIfAbsent(self.resultdir)
            except Error:
                raise ResultDirNotAccessible(ResultDirNotAccessible.__doc__ % self.resultdir)

    @traceLog()
    def _init(self, prebuild, do_log):

        self.state.start("chroot init")
        util.mkdirIfAbsent(self.basedir)
        mockgid = grp.getgrnam('mock').gr_gid
        os.chown(self.basedir, os.getuid(), mockgid)
        os.chmod(self.basedir, 0o2775)
        util.mkdirIfAbsent(self.make_chroot_path())
        self.plugins.call_hooks('mount_root')
        # intentionally we do not call bootstrap hook here - it does not have sense
        self._setup_nosync()
        self.chroot_was_initialized = self.chroot_is_initialized()
        self._setup_result_dir()
        getLog().info("calling preinit hooks")
        self.plugins.call_hooks('preinit')
        # intentionally we do not call bootstrap hook here - it does not have sense
        self.chroot_was_initialized = self.chroot_is_initialized()

        self._setup_dirs()
        if do_log:
            self._resetLogging()
        # /dev is later overwritten by systemd-nspawn, but we need this for
        # initial installation when chroot is empty
        self._setup_devices()

        self._setup_files()
        self.mounts.mountall_managed()

        # write out config details
        self.root_log.debug('rootdir = %s', self.make_chroot_path())
        self.root_log.debug('resultdir = %s', self.resultdir)

        self.pkg_manager.initialize()
        self._setup_resolver_config()
        if not self.chroot_was_initialized:
            self._setup_dbus_uuid()
            self._init_aux_files()
            if not util.USE_NSPAWN:
                self._setup_timezone()
            self._init_pkg_management()
            self._setup_files_postinstall()
            self._make_build_user()
            self._setup_build_dirs()
        elif prebuild:
            if 'age_check' in self.config['plugin_conf']['root_cache_opts'] and \
               not self.config['plugin_conf']['root_cache_opts']['age_check']:
                self._init_pkg_management()
            # Recreates build user to ensure the uid/gid are up to date with config
            # and there's no garbage left by previous build
            self._make_build_user()
            self._setup_build_dirs()
            if (self.config['online'] and self.config['update_before_build']
                    and self.config['clean']):
                update_state = "{0} update".format(self.pkg_manager.name)
                self.state.start(update_state)
                self.pkg_manager.update()
                self.state.finish(update_state)
        else:
            self._fixup_build_user()
            # Change owner of homdir tree if the root of it not owned
            # by the current user.
            home = self.make_chroot_path(self.homedir)
            if os.path.exists(home) and os.stat(home).st_uid != self.chrootuid:
                self.chown_home_dir()

        # mark the buildroot as initialized
        util.touch(self.make_chroot_path('.initialized'))

        # done with init
        self.plugins.call_hooks('postinit')
        # intentionally we do not call bootstrap hook here - it does not have sense

        self.mounts.mountall_user()

        self.state.finish("chroot init")

    def doChroot(self, command, shell=False, nosync=False, *args, **kargs):
        """execute given command in root"""
        self.nuke_rpm_db()
        env = dict(self.env)
        if nosync and self.nosync_path:
            env['LD_PRELOAD'] = self.nosync_path
        if util.USE_NSPAWN:
            if 'uid' not in kargs:
                kargs['uid'] = uid.getresuid()[1]
            if 'gid' not in kargs:
                kargs['gid'] = uid.getresgid()[1]
            if 'user' not in kargs:
                kargs['gid'] = pwd.getpwuid(kargs['uid'])[0]
            self.uid_manager.becomeUser(0, 0)
        result = util.do(command, chrootPath=self.make_chroot_path(),
                         env=env, shell=shell, *args, **kargs)
        if util.USE_NSPAWN:
            self.uid_manager.restorePrivs()
        return result

    @traceLog()
    def _copy_config(self, filename):
        etcdir = self.make_chroot_path('etc')
        conf_file = os.path.join(etcdir, filename)
        if os.path.exists(conf_file):
            os.remove(conf_file)
        orig_conf_file = os.path.join('/etc', filename)
        if os.path.exists(orig_conf_file):
            shutil.copy2(orig_conf_file, etcdir)
        else:
            self.root_log.warning("File %s not present. It is not copied into the chroot.", orig_conf_file)

    @traceLog()
    def _setup_resolver_config(self):
        if self.config['use_host_resolv'] and self.config['rpmbuild_networking']:
            self._copy_config('resolv.conf')
            self._copy_config('hosts')

    @traceLog()
    def _setup_dbus_uuid(self):
        machine_uuid = uuid.uuid4().hex
        dbus_uuid_path = self.make_chroot_path('etc', 'machine-id')
        symlink_path = self.make_chroot_path('var', 'lib', 'dbus', 'machine-id')
        with open(dbus_uuid_path, 'w') as uuid_file:
            uuid_file.write(machine_uuid)
            uuid_file.write('\n')
        if not os.path.exists(symlink_path):
            os.symlink("../../../etc/machine-id", symlink_path)

    @traceLog()
    def _setup_timezone(self):
        self._copy_config('localtime')

    @traceLog()
    def _init_pkg_management(self):
        update_state = '{0} install'.format(self.pkg_manager.name)
        self.state.start(update_state)
        if 'module_enable' in self.config and self.config['module_enable']:
            cmd = ['module', 'enable'] + self.config['module_enable']
            self.pkg_manager.init_install_output += self.pkg_manager.execute(*cmd, returnOutput=1)
        if 'module_install' in self.config and self.config['module_install']:
            cmd = ['module', 'install'] + self.config['module_install']
            self.pkg_manager.init_install_output += self.pkg_manager.execute(*cmd, returnOutput=1)

        cmd = self.config['chroot_setup_cmd']
        if cmd:
            if isinstance(cmd, util.basestring):


# multi-coloco

# input files - separate gwas summary files
# can specify region?
# each one with "SNP CHR BP BETA SE P"

# Then do:
# 1. merge summary stats
# 2. calc bfs
# 3. rbfs and ppas for each configuration

# test region
# chr15:78516053-80860978

import gzip
import itertools
import math
import os
import numpy as np
import string
import sys
import time
from scipy import stats

__version__ = '0.1'

def wakefield_bf(beta,se,p,w = 0.1):
  # calculate Wakefield bayes factor
  z = stats.norm.isf(p/2)
  r = w/(se**2 + w)
  bf = math.sqrt(1-r) * math.exp(z**2/2*r)
  return bf

def make_sigma(cor_mat,v,w):
  sigma = np.diag(v+w)
  for i in range(len(v)):
    for j in range(i+1,len(v)):
      c = cor_mat[i,j]
      sigma[i,j] = c * math.sqrt((v[i] + w[i]) * (v[j] + w[j]))
      sigma[j,i] = c * math.sqrt((v[i] + w[i]) * (v[j] + w[j]))
  return np.matrix(sigma)

def is_pos_def(A):
  #return np.all(np.linalg.eigvals(A) > 0)
  try:
    np.linalg.cholesky(A)
  except:
    return False
  return True

def prod(ls):
  prod = 1.0
  for i in ls:
    prod = prod * i
  return prod

def get_psd(A):
  diag = np.diag(A)
  detA = np.linalg.det(A)
  while not is_pos_def(A):
    #while detA <= 0:
    A = 0.999 * A
    for i in range(np.shape(A)[0]):
      A[i,i] = diag[i]
    #detA = np.linalg.det(A)
  return A

def is_num(x):
  try:
    float(x)
  except:
    return False
  return True

def gather_assocs(in_files,chrom,start,stop):
  # only include snps common to all traits
  # calculate wakefield bf
  # assoc_dict[snp][file1] = [snp,chrom,bp,beta,se,p,bf]
  n_files = len(in_files)
  assoc_dict = {}
  start = int(start)
  stop = int(stop)
  # summary stats and wakefield bfs
  for file in in_files:
    if file.split(".")[-1] == "gz":
      f = gzip.open(file,'r')
      header = gzip.open(file,'r').readline().split()
    else:
      f = open(file,'r')
      header = open(file,'r').readline().split()
    snp_i = header.index("SNP")
    chr_i = header.index("CHR")
    bp_i = header.index("BP")
    p_i = header.index("P")
    se_i = header.index("SE")
    if "OR" in header:
      is_odds = True
      or_i = header.index("OR")
    if "BETA" in header:
      is_odds = False
      beta_i = header.index("BETA")
    for i in f:
      line = i.split()
      # ignore NAs?
      if not is_num(line[p_i]) or not is_num(line[se_i]):
        continue
      if is_odds:
        if not is_num(line[or_i]):
          continue
      if not is_odds:
        if not is_num(line[beta_i]):
          continue
      if line[chr_i] != chrom:
        continue
      if int(line[bp_i]) >= start and int(line[bp_i]) <= stop:
        [snp,chrom,bp,se,p] = [line[j] for j in [snp_i,chr_i,bp_i,se_i,p_i]]
        if is_odds:
          beta = math.log(float(line[or_i]))
        else:
          beta = float(line[beta_i])
        bf = wakefield_bf(float(beta),float(se),float(p))
        if not snp in assoc_dict:
          assoc_dict[snp] = {file:[snp,chrom,bp,float(beta),float(se),float(p),bf]}
        else:
          assoc_dict[snp][file] = [snp,chrom,bp,float(beta),float(se),float(p),bf]
  bad_snps = [snp for snp in assoc_dict if len(assoc_dict[snp]) != n_files]
  for snp in bad_snps:
    del assoc_dict[snp]
  print "\nFound " + str(len(assoc_dict)) + " SNPs common to all " + str(n_files) + " traits in chr" + chrom + ":" + str(start) + "-" + str(stop)
  return assoc_dict

def adjust_bfs(assoc_dict,configs,n_files,in_files,overlap):
  # adjusted bfs
  # gene correlation matrix
  # if assuming no overlap, cor_mat is identity matrix
  cor_mat = np.identity(n_files)
  if overlap:
    for i in range(n_files):
      for j in range(i+1,n_files):
        beta1 = [assoc_dict[snp][in_files[i]][3] for snp in assoc_dict]
        beta2 = [assoc_dict[snp][in_files[j]][3] for snp in assoc_dict]
        cortest = stats.pearsonr(beta1,beta2)
        if cortest[1] >= 0.01:
          cor_mat[i,j] = 0.0
          cor_mat[j,i] = 0.0
        else:
          cor_mat[i,j] = cortest[0]
          cor_mat[j,i] = cortest[0]
    cor_mat = np.matrix(cor_mat)
  # configs - different adj_bf depending on configs
  # traits abcde...
  # get adj_bf for each config combo
  adj_bf_dict = {config:{snp:0.0 for snp in assoc_dict} for config in configs}
  for snp in assoc_dict:
    v = np.array([assoc_dict[snp][in_file][4]**2 for in_file in in_files])
    betas = np.array([assoc_dict[snp][in_file][3] for in_file in in_files])
    ses = np.array([assoc_dict[snp][in_file][4] for in_file in in_files])
    means = np.array([0.0 for i in range(n_files)])
    w_0 = np.array([0.0 for i in range(n_files)])
    sigma_h0 = make_sigma(cor_mat,v,w_0)
    if not is_pos_def(sigma_h0):
      sigma_h0 = get_psd(sigma_h0)
    #
    for config in configs:
      w = [0.0 for i in range(n_files)]
      for i in config:
        w[string.ascii_lowercase.index(i)] = 0.1
      sigma_h1 = make_sigma(cor_mat,v,w)
      if not is_pos_def(sigma_h1):
        sigma_h1 = get_psd(sigma_h1)
      # in case denominator is zero
      mvn_numer = stats.multivariate_normal.pdf(betas, means, sigma_h1)
      mvn_denom = stats.multivariate_normal.pdf(betas,means,sigma_h0)
      if mvn_denom == 0:
        adj_bf = 1.0
      else:
        adj_bf = mvn_numer/mvn_denom
      #adj_bf = stats.multivariate_normal.pdf(betas, means, sigma_h1) / stats.multivariate_normal.pdf(betas,means,sigma_h0)
      adj_bf_dict[config][snp] = adj_bf
  return adj_bf_dict

def get_final_configs(configs,n_files):
  # list of all possible configurations
  combs = configs[:]
  for iter in range(1,n_files+1):
    #print iter
    new_configs = []
    for config in configs:
      if len(config.split(",")) != iter:
        continue
      for comb in combs:
        overlap = 0
        raw_comb = "".join(comb.split(","))
        raw_config = "".join(config.split(","))
        # 1. check to see no numbers in comb and config overlap
        for i in raw_comb:
          if i in raw_config:
            overlap = 1
            break
        # 2. check that length is less than n_genes
        if len(raw_comb + raw_config) > n_files+1:
          overlap = 1
          continue
        # check that config hasn't been seen
        new_config = sorted([comb,config])
        if not new_config in new_configs and overlap == 0:
          new_configs.append(new_config)
    configs = configs + [",".join(i) for i in new_configs]
  final_configs = []
  for i in configs:
    config = ",".join([k for k in sorted([j for j in i.split(",")])])
    if not config in final_configs:
      final_configs.append(config)
  return final_configs

def do_moloco(adj_bf_dict,final_configs,configs,priors):
  single_bfs = {config:0.0 for config in configs}
  for config in configs:
    n_snps = len(adj_bf_dict[config])
    prior = priors[len(config)-1]
    trait_bfs = []
    for snp in adj_bf_dict[config]:
      trait_bfs.append(adj_bf_dict[config][snp])
    single_bfs[config] = priors[len(config)-1] * sum(trait_bfs)

  config_dict = {i:0.0 for i in final_configs}
  for config in final_configs:
    if config in single_bfs:
      config_dict[config] = single_bfs[config]
      continue
    # left side
    left_trait_bfs = 1.0
    for i in config.split(","):
      left_trait_bfs = left_trait_bfs * single_bfs[i]
    # right side
    prior = prod([priors[len(j)-1] for j in config.split(",")]) / priors[len("".join(config).replace(",",""))-1]
    right_trait_bfs = prod([priors[len(j)-1] for j in config.split(",")]) * single_bfs["".join(sorted(config.replace(",","")))]
    config_bf = left_trait_bfs - right_trait_bfs
    config_dict[config] = config_bf

  # add config where nothing is associated
  config_ppas = {config:config_dict[config]/(1.0 + sum(config_dict.values())) for config in final_configs}
  moloco_stats = {}
  for config in config_dict:
    moloco_stats[config] = [config_dict[config],config_ppas[config]]
    moloco_stats["zero"] = [1.0,1.0 - sum(config_ppas.values())]
  return moloco_stats

def moloc_iter(in_files,chrom,start,stop,priors,out_file,overlap,print_sum = False):
  # calc bf and consolidate sum stats
  n_files = len(in_files)
  assoc_dict = gather_assocs(in_files,chrom,start,stop)
  if len(assoc_dict) == 0:
    "Moving on to the next region"
    return True
  # single configs
  num2alpha = dict(zip(range(0, 26), string.ascii_lowercase))
  configs = [num2alpha[i] for i in range(n_files)]
  for i in range(1,n_files):
    n_causal = i + 1
    sel = "".join([num2alpha[k] for k in [j for j in range(n_files)]])
    configs = configs + ["".join(j) for j in list(itertools.combinations(sel,n_causal))]

  adj_bf_dict = adjust_bfs(assoc_dict,configs,n_files,in_files,overlap)
  # all configs
  final_configs = get_final_configs(configs,n_files)
  moloco = do_moloco(adj_bf_dict,final_configs,configs,priors)

  out_path = out_file + "." + chrom + "." + str(start) + "." + str(stop) + ".moloco"
  print "Great success! Writing results to " + out_path
  write_out = open(out_path,'wa')
  print >>write_out, "config logBF PP"
  print >>write_out, "0 0 " + str(moloco["zero"][1])
  for config in final_configs:
    print >>write_out, config + " " + str(math.log(moloco[config][0])) + " " + str(moloco[config][1])

  if print_sum:
    # Print summary stats
    print "Writting summary statistics to " + out_file + "." + chrom + "." + str(start) + "." + str(stop) + ".stats"
    write_sum = open(out_file + "." + chrom + "." + str(start) + "." + str(stop) + ".stats",'wa')
    print >>write_sum, "SNP CHR BP " + " ".join(["logBF_" + config for config in configs])
    for snp in assoc_dict:
      trait = assoc_dict[snp].keys()[0]
      print >>write_sum, assoc_dict[snp][trait][0] + " " + assoc_dict[snp][trait][1] + " " + assoc_dict[snp][trait][2],
      for config in configs:
        print >>write_sum, str(math.log(adj_bf_dict[config][snp])),
      print >>write_sum, ""
    write_sum.close()

  print ""
  # print out colocalization posteriors for each trait
  for i in string.ascii_lowercase[:n_files]:
    pp = 0.0
    for j in moloco:
      for k in j.split(","):
        if len(k) > 1 and i in k:
          pp = pp + moloco[j][1]
    print "Probability that trait \"" + i + "\" colocalizes with at least one other trait = " + str(pp)

  print ""
  return moloco

def main():
  # chrom = 15
  # start = 78516053
  # stop = 80860978
  #in_files = "lung_cancer.gwax.assoc.gz,bronchitis.gwax.assoc.gz,heart_disease.gwax.assoc.gz"
  start_time = time.time()
  print "\nMOLOCO"
  print "Version 0.1"
  print "Direct complaints to: jliu@nygenome.org\n"
  args = sys.argv[1:]
  do_bed = False
  overlap = True
  print_sum = False
  for i in range(len(args)):
    if args[i] == "--stats":
      in_files = args[i+1]
      print "Summary statistics: " + in_files


# CHILDES XML Corpus Reader

# Copyright (C) 2001-2015 NLTK Project
# Author: Tomonori Nagano <tnagano@gc.cuny.edu>
#         Alexis Dimitriadis <A.Dimitriadis@uu.nl>
# URL: <http://nltk.org/>
# For license information, see LICENSE.TXT

"""
Corpus reader for the XML version of the CHILDES corpus.
"""
from __future__ import print_function

__docformat__ = 'epytext en'

import re
from collections import defaultdict

from nltk.util import flatten
from nltk.compat import string_types

from nltk.corpus.reader.util import concat
from nltk.corpus.reader.xmldocs import XMLCorpusReader, ElementTree

# to resolve the namespace issue
NS = 'http://www.talkbank.org/ns/talkbank'

class CHILDESCorpusReader(XMLCorpusReader):
    """
    Corpus reader for the XML version of the CHILDES corpus.
    The CHILDES corpus is available at ``http://childes.psy.cmu.edu/``. The XML
    version of CHILDES is located at ``http://childes.psy.cmu.edu/data-xml/``.
    Copy the needed parts of the CHILDES XML corpus into the NLTK data directory
    (``nltk_data/corpora/CHILDES/``).

    For access to the file text use the usual nltk functions,
    ``words()``, ``sents()``, ``tagged_words()`` and ``tagged_sents()``.
    """
    def __init__(self, root, fileids, lazy=True):
        XMLCorpusReader.__init__(self, root, fileids)
        self._lazy = lazy

    def words(self, fileids=None, speaker='ALL', stem=False,
            relation=False, strip_space=True, replace=False):
        """
        :return: the given file(s) as a list of words
        :rtype: list(str)

        :param speaker: If specified, select specific speaker(s) defined
            in the corpus. Default is 'ALL' (all participants). Common choices
            are 'CHI' (the child), 'MOT' (mother), ['CHI','MOT'] (exclude
            researchers)
        :param stem: If true, then use word stems instead of word strings.
        :param relation: If true, then return tuples of (stem, index,
            dependent_index)
        :param strip_space: If true, then strip trailing spaces from word
            tokens. Otherwise, leave the spaces on the tokens.
        :param replace: If true, then use the replaced (intended) word instead
            of the original word (e.g., 'wat' will be replaced with 'watch')
        """
        sent=None
        pos=False
        return concat([self._get_words(fileid, speaker, sent, stem, relation,
            pos, strip_space, replace) for fileid in self.abspaths(fileids)])

    def tagged_words(self, fileids=None, speaker='ALL', stem=False,
            relation=False, strip_space=True, replace=False):
        """
        :return: the given file(s) as a list of tagged
            words and punctuation symbols, encoded as tuples
            ``(word,tag)``.
        :rtype: list(tuple(str,str))

        :param speaker: If specified, select specific speaker(s) defined
            in the corpus. Default is 'ALL' (all participants). Common choices
            are 'CHI' (the child), 'MOT' (mother), ['CHI','MOT'] (exclude
            researchers)
        :param stem: If true, then use word stems instead of word strings.
        :param relation: If true, then return tuples of (stem, index,
            dependent_index)
        :param strip_space: If true, then strip trailing spaces from word
            tokens. Otherwise, leave the spaces on the tokens.
        :param replace: If true, then use the replaced (intended) word instead
            of the original word (e.g., 'wat' will be replaced with 'watch')
        """
        sent=None
        pos=True
        return concat([self._get_words(fileid, speaker, sent, stem, relation,
            pos, strip_space, replace) for fileid in self.abspaths(fileids)])

    def sents(self, fileids=None, speaker='ALL', stem=False,
            relation=None, strip_space=True, replace=False):
        """
        :return: the given file(s) as a list of sentences or utterances, each
            encoded as a list of word strings.
        :rtype: list(list(str))

        :param speaker: If specified, select specific speaker(s) defined
            in the corpus. Default is 'ALL' (all participants). Common choices
            are 'CHI' (the child), 'MOT' (mother), ['CHI','MOT'] (exclude
            researchers)
        :param stem: If true, then use word stems instead of word strings.
        :param relation: If true, then return tuples of ``(str,pos,relation_list)``.
            If there is manually-annotated relation info, it will return
            tuples of ``(str,pos,test_relation_list,str,pos,gold_relation_list)``
        :param strip_space: If true, then strip trailing spaces from word
            tokens. Otherwise, leave the spaces on the tokens.
        :param replace: If true, then use the replaced (intended) word instead
            of the original word (e.g., 'wat' will be replaced with 'watch')
        """
        sent=True
        pos=False
        return concat([self._get_words(fileid, speaker, sent, stem, relation,
            pos, strip_space, replace) for fileid in self.abspaths(fileids)])

    def tagged_sents(self, fileids=None, speaker='ALL', stem=False,
            relation=None, strip_space=True, replace=False):
        """
        :return: the given file(s) as a list of
            sentences, each encoded as a list of ``(word,tag)`` tuples.
        :rtype: list(list(tuple(str,str)))

        :param speaker: If specified, select specific speaker(s) defined
            in the corpus. Default is 'ALL' (all participants). Common choices
            are 'CHI' (the child), 'MOT' (mother), ['CHI','MOT'] (exclude
            researchers)
        :param stem: If true, then use word stems instead of word strings.
        :param relation: If true, then return tuples of ``(str,pos,relation_list)``.
            If there is manually-annotated relation info, it will return
            tuples of ``(str,pos,test_relation_list,str,pos,gold_relation_list)``
        :param strip_space: If true, then strip trailing spaces from word
            tokens. Otherwise, leave the spaces on the tokens.
        :param replace: If true, then use the replaced (intended) word instead
            of the original word (e.g., 'wat' will be replaced with 'watch')
        """
        sent=True
        pos=True
        return concat([self._get_words(fileid, speaker, sent, stem, relation,
            pos, strip_space, replace) for fileid in self.abspaths(fileids)])

    def corpus(self, fileids=None):
        """
        :return: the given file(s) as a dict of ``(corpus_property_key, value)``
        :rtype: list(dict)
        """
        return [self._get_corpus(fileid) for fileid in self.abspaths(fileids)]

    def _get_corpus(self, fileid):
        results = dict()
        xmldoc = ElementTree.parse(fileid).getroot()
        for key, value in xmldoc.items():
            results[key] = value
        return results

    def participants(self, fileids=None):
        """
        :return: the given file(s) as a dict of
            ``(participant_property_key, value)``
        :rtype: list(dict)
        """
        return [self._get_participants(fileid)
                            for fileid in self.abspaths(fileids)]

    def _get_participants(self, fileid):
        # multidimensional dicts
        def dictOfDicts():
            return defaultdict(dictOfDicts)

        xmldoc = ElementTree.parse(fileid).getroot()
        # getting participants' data
        pat = dictOfDicts()
        for participant in xmldoc.findall('.//{%s}Participants/{%s}participant'
                                          % (NS,NS)):
            for (key,value) in participant.items():
                pat[participant.get('id')][key] = value
        return pat

    def age(self, fileids=None, speaker='CHI', month=False):
        """
        :return: the given file(s) as string or int
        :rtype: list or int

        :param month: If true, return months instead of year-month-date
        """
        return [self._get_age(fileid, speaker, month)
                for fileid in self.abspaths(fileids)]

    def _get_age(self, fileid, speaker, month):
        xmldoc = ElementTree.parse(fileid).getroot()
        for pat in xmldoc.findall('.//{%s}Participants/{%s}participant'
                                  % (NS,NS)):
            try:
                if pat.get('id') == speaker:
                    age = pat.get('age')
                    if month:
                        age = self.convert_age(age)
                    return age
            # some files don't have age data
            except (TypeError, AttributeError) as e:
                return None

    def convert_age(self, age_year):
        "Caclculate age in months from a string in CHILDES format"
        m = re.match("P(\d+)Y(\d+)M?(\d?\d?)D?",age_year)
        age_month = int(m.group(1))*12 + int(m.group(2))
        try:
            if int(m.group(3)) > 15:
                age_month += 1
        # some corpora don't have age information?
        except ValueError as e:
            pass
        return age_month

    def MLU(self, fileids=None, speaker='CHI'):
        """
        :return: the given file(s) as a floating number
        :rtype: list(float)
        """
        return [self._getMLU(fileid, speaker=speaker)
                for fileid in self.abspaths(fileids)]

    def _getMLU(self, fileid, speaker):
        sents = self._get_words(fileid, speaker=speaker, sent=True, stem=True,
                    relation=False, pos=True, strip_space=True, replace=True)
        results = []
        lastSent = []
        numFillers = 0
        sentDiscount = 0
        for sent in sents:
            posList = [pos for (word,pos) in sent]
            # if any part of the sentence is intelligible
            if any(pos == 'unk' for pos in posList):
                next
            # if the sentence is null
            elif sent == []:
                next
            # if the sentence is the same as the last sent
            elif sent == lastSent:
                next
            else:
                results.append([word for (word,pos) in sent])
                # count number of fillers
                if len(set(['co',None]).intersection(posList)) > 0:
                    numFillers += posList.count('co')
                    numFillers += posList.count(None)
                    sentDiscount += 1
            lastSent = sent
        try:
            thisWordList = flatten(results)
            # count number of morphemes
            # (e.g., 'read' = 1 morpheme but 'read-PAST' is 2 morphemes)
            numWords = float(len(flatten([word.split('-')
                                          for word in thisWordList]))) - numFillers
            numSents = float(len(results)) - sentDiscount
            mlu = numWords/numSents
        except ZeroDivisionError:
            mlu = 0
        # return {'mlu':mlu,'wordNum':numWords,'sentNum':numSents}
        return mlu

    def _get_words(self, fileid, speaker, sent, stem, relation, pos,
            strip_space, replace):
        if isinstance(speaker, string_types) and speaker != 'ALL':  # ensure we have a list of speakers
            speaker = [ speaker ]
        xmldoc = ElementTree.parse(fileid).getroot()
        # processing each xml doc
        results = []
        for xmlsent in xmldoc.findall('.//{%s}u' % NS):
            sents = []
            # select speakers
            if speaker == 'ALL' or xmlsent.get('who') in speaker:
                for xmlword in xmlsent.findall('.//{%s}w' % NS):
                    infl = None ; suffixStem = None
                    # getting replaced words
                    if replace and xmlsent.find('.//{%s}w/{%s}replacement'
                                                % (NS,NS)):
                        xmlword = xmlsent.find('.//{%s}w/{%s}replacement/{%s}w'
                                               % (NS,NS,NS))
                    elif replace and xmlsent.find('.//{%s}w/{%s}wk' % (NS,NS)):
                        xmlword = xmlsent.find('.//{%s}w/{%s}wk' % (NS,NS))
                    # get text
                    if xmlword.text:
                        word = xmlword.text
                    else:
                        word = ''
                    # strip tailing space
                    if strip_space:
                        word = word.strip()
                    # stem
                    if relation or stem:
                        try:
                            xmlstem = xmlword.find('.//{%s}stem' % NS)
                            word = xmlstem.text
                        except AttributeError as e:
                            pass
                        # if there is an inflection
                        try:
                            xmlinfl = xmlword.find('.//{%s}mor/{%s}mw/{%s}mk'
                                                   % (NS,NS,NS))
                            word += '-' + xmlinfl.text
                        except:
                            pass
                        # if there is a suffix
                        try:
                            xmlsuffix = xmlword.find('.//{%s}mor/{%s}mor-post/{%s}mw/{%s}stem'
                                                     % (NS,NS,NS,NS))
                            suffixStem = xmlsuffix.text
                        except AttributeError:
                            suffixStem = ""
                    # pos
                    if relation or pos:
                        try:
                            xmlpos = xmlword.findall(".//{%s}c" % NS)
                            xmlpos2 = xmlword.findall(".//{%s}s" % NS)


"""
Handlers for video module.

StudentViewHandlers are handlers for video module instance.
StudioViewHandlers are handlers for video descriptor instance.
"""

import json
import logging
import math

import six
from django.core.files.base import ContentFile
from django.utils.timezone import now
from edxval.api import create_external_video, create_or_update_video_transcript, delete_video_transcript
from opaque_keys.edx.locator import CourseLocator
from webob import Response
from xblock.core import XBlock
from xblock.exceptions import JsonHandlerError

from xmodule.exceptions import NotFoundError
from xmodule.fields import RelativeTime

from .transcripts_utils import (
    Transcript,
    TranscriptException,
    TranscriptsGenerationException,
    clean_video_id,
    generate_sjson_for_all_speeds,
    get_html5_ids,
    get_or_create_sjson,
    get_transcript,
    get_transcript_from_contentstore,
    remove_subs_from_store,
    subs_filename,
    youtube_speed_dict
)

log = logging.getLogger(__name__)

# Disable no-member warning:
# pylint: disable=no-member

def to_boolean(value):
    """
    Convert a value from a GET or POST request parameter to a bool
    """
    if isinstance(value, six.binary_type):
        value = value.decode('ascii', errors='replace')
    if isinstance(value, six.text_type):
        return value.lower() == 'true'
    else:
        return bool(value)

class VideoStudentViewHandlers(object):
    """
    Handlers for video module instance.
    """

    def handle_ajax(self, dispatch, data):
        """
        Update values of xfields, that were changed by student.
        """
        accepted_keys = [
            'speed', 'auto_advance', 'saved_video_position', 'transcript_language',
            'transcript_download_format', 'youtube_is_available',
            'bumper_last_view_date', 'bumper_do_not_show_again'
        ]

        conversions = {
            'speed': json.loads,
            'auto_advance': json.loads,
            'saved_video_position': RelativeTime.isotime_to_timedelta,
            'youtube_is_available': json.loads,
            'bumper_last_view_date': to_boolean,
            'bumper_do_not_show_again': to_boolean,
        }

        if dispatch == 'save_user_state':
            for key in data:
                if key in accepted_keys:
                    if key in conversions:
                        value = conversions[key](data[key])
                    else:
                        value = data[key]

                    if key == 'bumper_last_view_date':
                        value = now()

                    if key == 'speed' and math.isnan(value):
                        message = u"Invalid speed value {}, must be a float.".format(value)
                        log.warning(message)
                        return json.dumps({'success': False, 'error': message})

                    setattr(self, key, value)

                    if key == 'speed':
                        self.global_speed = self.speed  # lint-amnesty, pylint: disable=attribute-defined-outside-init

            return json.dumps({'success': True})

        log.debug(u"GET {0}".format(data))
        log.debug(u"DISPATCH {0}".format(dispatch))

        raise NotFoundError('Unexpected dispatch type')

    def translation(self, youtube_id, transcripts):
        """
        This is called to get transcript file for specific language.

        youtube_id: str: must be one of youtube_ids or None if HTML video
        transcripts (dict): A dict with all transcripts and a sub.

        Logic flow:

        If youtube_id doesn't exist, we have a video in HTML5 mode. Otherwise,
        video in Youtube or Flash modes.

        if youtube:
            If english -> give back youtube_id subtitles:
                Return what we have in contentstore for given youtube_id.
            If non-english:
                a) extract youtube_id from srt file name.
                b) try to find sjson by youtube_id and return if successful.
                c) generate sjson from srt for all youtube speeds.
        if non-youtube:
            If english -> give back `sub` subtitles:
                Return what we have in contentstore for given subs_if that is stored in self.sub.
            If non-english:
                a) try to find previously generated sjson.
                b) otherwise generate sjson from srt and return it.

        Filenames naming:
            en: subs_videoid.srt.sjson
            non_en: uk_subs_videoid.srt.sjson

        Raises:
            NotFoundError if for 'en' subtitles no asset is uploaded.
            NotFoundError if youtube_id does not exist / invalid youtube_id
        """
        sub, other_lang = transcripts["sub"], transcripts["transcripts"]
        if youtube_id:
            # Youtube case:
            if self.transcript_language == 'en':
                return Transcript.asset(self.location, youtube_id).data

            youtube_ids = youtube_speed_dict(self)
            if youtube_id not in youtube_ids:
                log.info("Youtube_id %s does not exist", youtube_id)
                raise NotFoundError

            try:
                sjson_transcript = Transcript.asset(self.location, youtube_id, self.transcript_language).data
            except NotFoundError:
                log.info("Can't find content in storage for %s transcript: generating.", youtube_id)
                generate_sjson_for_all_speeds(
                    self,
                    other_lang[self.transcript_language],
                    {speed: youtube_id for youtube_id, speed in six.iteritems(youtube_ids)},
                    self.transcript_language
                )
                sjson_transcript = Transcript.asset(self.location, youtube_id, self.transcript_language).data

            return sjson_transcript
        else:
            # HTML5 case
            if self.transcript_language == 'en':
                if '.srt' not in sub:  # not bumper case
                    return Transcript.asset(self.location, sub).data
                try:
                    return get_or_create_sjson(self, {'en': sub})
                except TranscriptException:
                    pass  # to raise NotFoundError and try to get data in get_static_transcript
            elif other_lang:
                return get_or_create_sjson(self, other_lang)

        raise NotFoundError

    def get_static_transcript(self, request, transcripts):
        """
        Courses that are imported with the --nostatic flag do not show
        transcripts/captions properly even if those captions are stored inside
        their static folder. This adds a last resort method of redirecting to
        the static asset path of the course if the transcript can't be found
        inside the contentstore and the course has the static_asset_path field
        set.

        transcripts (dict): A dict with all transcripts and a sub.
        """
        response = Response(status=404)
        # Only do redirect for English
        if not self.transcript_language == 'en':
            return response

        # If this video lives in library, the code below is not relevant and will error.
        if not isinstance(self.course_id, CourseLocator):
            return response

        video_id = request.GET.get('videoId', None)
        if video_id:
            transcript_name = video_id
        else:
            transcript_name = transcripts["sub"]

        if transcript_name:
            # Get the asset path for course
            asset_path = None
            course = self.runtime.modulestore.get_course(self.course_id)
            if course.static_asset_path:
                asset_path = course.static_asset_path
            else:
                # It seems static_asset_path is not set in any XMLModuleStore courses.
                asset_path = getattr(course, 'data_dir', '')

            if asset_path:
                response = Response(
                    status=307,
                    location='/static/{0}/{1}'.format(
                        asset_path,
                        subs_filename(transcript_name, self.transcript_language)
                    )
                )
        return response

    @XBlock.json_handler
    def publish_completion(self, data, dispatch):  # pylint: disable=unused-argument
        """
        Entry point for completion for student_view.

        Parameters:
            data: JSON dict:
                key: "completion"
                value: float in range [0.0, 1.0]

            dispatch: Ignored.
        Return value: JSON response (200 on success, 400 for malformed data)
        """
        completion_service = self.runtime.service(self, 'completion')
        if completion_service is None:  # lint-amnesty, pylint: disable=no-else-raise
            raise JsonHandlerError(500, u"No completion service found")
        elif not completion_service.completion_tracking_enabled():
            raise JsonHandlerError(404, u"Completion tracking is not enabled and API calls are unexpected")
        if not isinstance(data['completion'], (int, float)):  # lint-amnesty, pylint: disable=no-else-raise
            message = u"Invalid completion value {}. Must be a float in range [0.0, 1.0]"
            raise JsonHandlerError(400, message.format(data['completion']))
        elif not 0.0 <= data['completion'] <= 1.0:
            message = u"Invalid completion value {}. Must be in range [0.0, 1.0]"
            raise JsonHandlerError(400, message.format(data['completion']))
        self.runtime.publish(self, "completion", data)
        return {"result": "ok"}

    @staticmethod
    def make_transcript_http_response(content, filename, language, content_type, add_attachment_header=True):
        """
        Construct `Response` object.

        Arguments:
            content (unicode): transcript content
            filename (unicode): transcript filename
            language (unicode): transcript language
            mimetype (unicode): transcript content type
            add_attachment_header (bool): whether to add attachment header or not
        """
        headerlist = [
            ('Content-Language', language),
        ]

        if add_attachment_header:
            headerlist.append(
                (
                    'Content-Disposition',
                    'attachment; filename="{}"'.format(filename.encode('utf-8') if six.PY2 else filename)
                )
            )

        response = Response(
            content,
            headerlist=headerlist,
            charset='utf8'
        )
        response.content_type = content_type

        return response

    @XBlock.handler
    def transcript(self, request, dispatch):
        """
        Entry point for transcript handlers for student_view.

        Request GET contains:
            (optional) `videoId` for `translation` dispatch.
            `is_bumper=1` flag for bumper case.

        Dispatches, (HTTP GET):
            /translation/[language_id]
            /download
            /available_translations/

        Explanations:
            `download`: returns SRT or TXT file.
            `translation`: depends on HTTP methods:
                    Provide translation for requested language, SJSON format is sent back on success,
                    Proper language_id should be in url.
            `available_translations`:
                    Returns list of languages, for which transcript files exist.
                    For 'en' check if SJSON exists. For non-`en` check if SRT file exists.
        """
        is_bumper = request.GET.get('is_bumper', False)
        transcripts = self.get_transcripts_info(is_bumper)

        if dispatch.startswith('translation'):
            language = dispatch.replace('translation', '').strip('/')

            if not language:
                log.info("Invalid /translation request: no language.")
                return Response(status=400)

            if language not in ['en'] + list(transcripts["transcripts"].keys()):
                log.info("Video: transcript facilities are not available for given language.")
                return Response(status=404)

            if language != self.transcript_language:  # lint-amnesty, pylint: disable=access-member-before-definition
                self.transcript_language = language  # lint-amnesty, pylint: disable=attribute-defined-outside-init

            try:
                if is_bumper:
                    content, filename, mimetype = get_transcript_from_contentstore(
                        self,
                        self.transcript_language,
                        Transcript.SJSON,
                        transcripts
                    )
                else:
                    content, filename, mimetype = get_transcript(
                        self,
                        lang=self.transcript_language,
                        output_format=Transcript.SJSON,
                        youtube_id=request.GET.get('videoId'),
                    )

                response = self.make_transcript_http_response(
                    content,
                    filename,
                    self.transcript_language,
                    mimetype,
                    add_attachment_header=False
                )
            except NotFoundError:
                log.exception('[Translation Dispatch] %s', self.location)
                response = self.get_static_transcript(request, transcripts)

        elif dispatch == 'download':
            lang = request.GET.get('lang', None)

            try:
                content, filename, mimetype = get_transcript(self, lang, output_format=self.transcript_download_format)
            except NotFoundError:
                return Response(status=404)

            response = self.make_transcript_http_response(
                content,
                filename,
                self.transcript_language,
                mimetype
            )
        elif dispatch.startswith('available_translations'):
            available_translations = self.available_translations(
                transcripts,
                verify_assets=True,
                is_bumper=is_bumper
            )
            if available_translations:
                response = Response(json.dumps(available_translations))
                response.content_type = 'application/json'
            else:
                response = Response(status=404)
        else:  # unknown dispatch
            log.debug("Dispatch is not allowed")
            response = Response(status=404)

        return response

    @XBlock.handler
    def student_view_user_state(self, request, suffix=''):  # lint-amnesty, pylint: disable=unused-argument
        """
        Endpoint to get user-specific state, like current position and playback speed,
        without rendering the full student_view HTML. This is similar to student_view_state,
        but that one cannot contain user-specific info.
        """
        view_state = self.student_view_data()
        view_state.update({
            "saved_video_position": self.saved_video_position.total_seconds(),
            "speed": self.speed,
        })
        return Response(
            json.dumps(view_state),
            content_type='application/json',
            charset='UTF-8'
        )

    @XBlock.handler
    def yt_video_metadata(self, request, suffix=''):  # lint-amnesty, pylint: disable=unused-argument
        """
        Endpoint to get YouTube metadata.
        This handler is only used in the Blockstore-based runtime. The old
        runtime uses a similar REST API that's not an XBlock handler.
        """
        from lms.djangoapps.courseware.views.views import load_metadata_from_youtube
        if not self.youtube_id_1_0:
            # TODO: more informational response to explain that yt_video_metadata not supported for non-youtube videos.
            return Response('{}', status=400)

        metadata, status_code = load_metadata_from_youtube(video_id=self.youtube_id_1_0, request=request)
        response = Response(json.dumps(metadata), status=status_code)
        response.content_type = 'application/json'
        return response

class VideoStudioViewHandlers(object):
    """
    Handlers for Studio view.
    """
    def validate_transcript_upload_data(self, data):
        """
        Validates video transcript file.
        Arguments:
            data: Transcript data to be validated.
        Returns:
            None or String
            If there is error returns error message otherwise None.
        """
        error = None
        _ = self.runtime.service(self, "i18n").ugettext
        # Validate the must have attributes - this error is unlikely to be faced by common users.
        must_have_attrs = ['edx_video_id', 'language_code', 'new_language_code']
        missing = [attr for attr in must_have_attrs if attr not in data]

        # Get available transcript languages.
        transcripts = self.get_transcripts_info()
        available_translations = self.available_translations(transcripts, verify_assets=True)

        if missing:
            error = _(u'The following parameters are required: {missing}.').format(missing=', '.join(missing))
        elif (
            data['language_code'] != data['new_language_code'] and data['new_language_code'] in available_translations
        ):
            error = _(u'A transcript with the "{language_code}" language code already exists.'.format(  # lint-amnesty, pylint: disable=translation-of-non-string
                language_code=data['new_language_code']
            ))
        elif 'file' not in data:
            error = _(u'A transcript file is required.')

        return error

    @XBlock.handler
    def studio_transcript(self, request, dispatch):
        """
        Entry point for Studio transcript handlers.

        Dispatches:
            /translation/[language_id] - language_id sould be in url.

        `translation` dispatch support following HTTP methods:
            `POST`:
                Upload srt file. Check possibility of generation of proper sjson files.
                For now, it works only for self.transcripts, not for `en`.
                Do not update self.transcripts, as fields are updated on save in Studio.
            `GET:
                Return filename from storage. SRT format is sent back on success. Filename should be in GET dict.

        We raise all exceptions right in Studio:
            NotFoundError:
                Video or asset was deleted from module/contentstore, but request came later.
                Seems impossible to be raised. module_render.py catches NotFoundErrors from here.

            /translation POST:
                TypeError:
                    Unjsonable filename or content.
                TranscriptsGenerationException, TranscriptException:
                    no SRT extension or not parse-able by PySRT
                UnicodeDecodeError: non-UTF8 uploaded file content encoding.
        """
        _ = self.runtime.service(self, "i18n").ugettext

        if dispatch.startswith('translation'):

            if request.method == 'POST':
                error = self.validate_transcript_upload_data(data=request.POST)
                if error:
                    response = Response(json={'error': error}, status=400)
                else:
                    edx_video_id = clean_video_id(request.POST['edx_video_id'])
                    language_code = request.POST['language_code']
                    new_language_code = request.POST['new_language_code']
                    transcript_file = request.POST['file'].file

                    if not edx_video_id:
                        # Back-populate the video ID for an external video.
                        # pylint: disable=attribute-defined-outside-init
                        self.edx_video_id = edx_video_id = create_external_video(display_name=u'external video')

                    try:
                        # Convert SRT transcript into an SJSON format
                        # and upload it to S3.
                        sjson_subs = Transcript.convert(
                            content=transcript_file.read().decode('utf-8'),
                            input_format=Transcript.SRT,
                            output_format=Transcript.SJSON
                        ).encode()
                        create_or_update_video_transcript(
                            video_id=edx_video_id,
                            language_code=language_code,
                            metadata={
                                'file_format': Transcript.SJSON,
                                'language_code': new_language_code
                            },
                            file_data=ContentFile(sjson_subs),
                        )
                        payload = {
                            'edx_video_id': edx_video_id,
                            'language_code': new_language_code
                        }
                        response = Response(json.dumps(payload), status=201)
                    except (TranscriptsGenerationException, UnicodeDecodeError):
                        response = Response(
                            json={
                                'error': _(
                                    u'There is a problem with this transcript file. Try to upload a different file.'
                                )
                            },
                            status=400
                        )
            elif request.method == 'DELETE':
                request_data = request.json

                if 'lang' not in request_data or 'edx_video_id' not in request_data:
                    return Response(status=400)

                language = request_data['lang']
                edx_video_id = clean_video_id(request_data['edx_video_id'])

                if edx_video_id:
                    delete_video_transcript(video_id=edx_video_id, language_code=language)

                if language == u'en':
                    # remove any transcript file from content store for the video ids
                    possible_sub_ids = [
                        self.sub,  # pylint: disable=access-member-before-definition
                        self.youtube_id_1_0
                    ] + get_html5_ids(self.html5_sources)
                    for sub_id in possible_sub_ids:
                        remove_subs_from_store(sub_id, self, language)

                    # update metadata as `en` can also be present in `transcripts` field
                    remove_subs_from_store(self.transcripts.pop(language, None), self, language)

                    # also empty `sub` field
                    self.sub = ''  # pylint: disable=attribute-defined-outside-init
                else:
                    remove_subs_from_store(self.transcripts.pop(language, None), self, language)

                return Response(status=200)

            elif request.method == 'GET':


# This file is part of Indico.
# Copyright (C) 2002 - 2016 European Organization for Nuclear Research (CERN).
#
# Indico is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License as
# published by the Free Software Foundation; either version 3 of the
# License, or (at your option) any later version.
#
# Indico is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Indico; if not, see <http://www.gnu.org/licenses/>.

# TODO: make --rb-zodb-uri optional and assume same DB in that case?
from collections import defaultdict
import os
import re

import time
from datetime import datetime, timedelta
from itertools import ifilter
from operator import attrgetter

import click
from babel import dates

from indico.core.db.sqlalchemy import db
from indico.modules.rb import settings as rb_settings
from indico.modules.rb.models.aspects import Aspect
from indico.modules.rb.models.blocked_rooms import BlockedRoom
from indico.modules.rb.models.blocking_principals import BlockingPrincipal
from indico.modules.rb.models.blockings import Blocking
from indico.modules.rb.models.locations import Location
from indico.modules.rb.models.photos import Photo
from indico.modules.rb.models.reservation_edit_logs import ReservationEditLog
from indico.modules.rb.models.reservation_occurrences import ReservationOccurrence
from indico.modules.rb.models.reservations import RepeatMapping, Reservation
from indico.modules.rb.models.room_attributes import RoomAttributeAssociation, RoomAttribute
from indico.modules.rb.models.room_bookable_hours import BookableHours
from indico.modules.rb.models.equipment import EquipmentType
from indico.modules.rb.models.room_nonbookable_periods import NonBookablePeriod
from indico.modules.rb.models.rooms import Room
from indico.util.console import cformat, verbose_iterator
from indico.util.date_time import as_utc
from indico.util.string import safe_upper, is_valid_mail
from indico_zodbimport import Importer, convert_to_unicode, convert_principal_list
from indico_zodbimport.util import UnbreakingDB, get_storage

french_month_names = [(str(i), name[:3].encode('utf-8').lower())
                      for i, name in dates.get_month_names(locale='fr_FR').iteritems()]

attribute_map = {
    'Simba List': 'Manager Group',
    'Booking Simba List': 'Allowed Booking Group'
}

def generate_name(old_room):
    return '{}-{}-{}'.format(old_room.building, old_room.floor, old_room.roomNr)

def get_canonical_name_of(old_room):
    return '{}-{}'.format(old_room._locationName, generate_name(old_room))

def get_room_id(guid):
    return int(guid.split('|')[1].strip())

def parse_dt_string(value):
    try:
        return datetime.strptime(value, '%d %b %Y %H:%M')
    except ValueError:
        # French month name
        for num, name in french_month_names:
            if name in value:
                value = value.lower().replace(name, num)
                break
        return datetime.strptime(value, '%d %m %Y %H:%M')

def utc_to_local(dt):
    assert dt.tzinfo is None
    return dt - timedelta(seconds=time.altzone)

# noinspection PyArgumentList
class RoomBookingImporter(Importer):
    def __init__(self, **kwargs):
        self.rb_zodb_uri = kwargs.pop('rb_zodb_uri')
        self.no_merged_avatars = kwargs.pop('no_merged_avatars')
        self.photo_path = kwargs.pop('photo_path')
        self.rb_root = None
        self.merged_avatars = {}
        super(RoomBookingImporter, self).__init__(**kwargs)

    @staticmethod
    def decorate_command(command):
        command = click.option('--rb-zodb-uri', required=True, help="ZODB URI for the room booking database")(command)
        command = click.option('--photo-path', type=click.Path(exists=True, file_okay=False),
                               help="path to the folder containing room photos")(command)
        command = click.option('--no-merged-avatars', default=False, is_flag=True,
                               help="speed up migration by not checking for merged avatars")(command)
        return command

    def connect_zodb(self):
        super(RoomBookingImporter, self).connect_zodb()
        self.rb_root = UnbreakingDB(get_storage(self.rb_zodb_uri)).open().root()

    def has_data(self):
        return bool(Location.find().count())

    def migrate(self):
        if not self.no_merged_avatars:
            self.find_merged_avatars()
        self.migrate_settings()
        self.migrate_locations()
        self.migrate_rooms()
        self.migrate_blockings()
        self.migrate_reservations()
        self.fix_sequences(schema='roombooking')

    def find_merged_avatars(self):
        print cformat('%{white!}checking for merged avatars')
        for avatar in verbose_iterator(self.zodb_root['avatars'].itervalues(), len(self.zodb_root['avatars']),
                                       attrgetter('id'), lambda av: '{}, {}'.format(safe_upper(av.surName), av.name)):
            for merged_avatar in getattr(avatar, '_mergeFrom', []):
                self.merged_avatars[merged_avatar.getId()] = avatar.getId()

    def migrate_settings(self):
        print cformat('%{white!}migrating settings')
        rb_settings.delete_all()
        opts = self.zodb_root['plugins']['RoomBooking']._PluginBase__options

        # Admins & authorized users/groups
        rb_settings.set('authorized_principals', convert_principal_list(opts['AuthorisedUsersGroups']))
        rb_settings.set('admin_principals', convert_principal_list(opts['Managers']))
        # Assistance emails
        emails = [email for email in opts['assistanceNotificationEmails']._PluginOption__value
                  if is_valid_mail(email, False)]
        rb_settings.set('assistance_emails', emails)
        # Simple settings
        rb_settings.set('notification_hour', opts['notificationHour']._PluginOption__value)
        rb_settings.set('notification_before_days', opts['notificationBefore']._PluginOption__value)
        db.session.commit()

    def migrate_locations(self):
        print cformat('%{white!}migrating locations')
        default_location_name = self.zodb_root['DefaultRoomBookingLocation']
        custom_attributes_dict = self.rb_root['CustomAttributesList']

        for old_location in self.zodb_root['RoomBookingLocationList']:
            # create location
            l = Location(
                name=convert_to_unicode(old_location.friendlyName),
                is_default=(old_location.friendlyName == default_location_name)
            )

            print cformat('- %{cyan}{}').format(l.name)

            # add aspects
            for old_aspect in old_location._aspects.values():
                a = Aspect(
                    name=convert_to_unicode(old_aspect.name),
                    center_latitude=old_aspect.centerLatitude,
                    center_longitude=old_aspect.centerLongitude,
                    zoom_level=old_aspect.zoomLevel,
                    top_left_latitude=old_aspect.topLeftLatitude,
                    top_left_longitude=old_aspect.topLeftLongitude,
                    bottom_right_latitude=old_aspect.bottomRightLatitude,
                    bottom_right_longitude=old_aspect.bottomRightLongitude
                )

                print cformat('  %{blue!}Aspect:%{reset} {}').format(a.name)

                l.aspects.append(a)
                if old_aspect.defaultOnStartup:
                    l.default_aspect = a

            # add custom attributes
            for ca in custom_attributes_dict.get(l.name, []):
                if ca['type'] != 'str':
                    raise RuntimeError('Non-str custom attributes are unsupported: {}'.format(ca))
                attr_name = attribute_map.get(ca['name'], ca['name'])
                attr = RoomAttribute(name=attr_name.replace(' ', '-').lower(), title=attr_name, type=ca['type'],
                                     is_required=ca['required'], is_hidden=ca['hidden'])
                l.attributes.append(attr)
                print cformat('  %{blue!}Attribute:%{reset} {}').format(attr.title)

            # add new created location
            db.session.add(l)
            print
            print
        db.session.commit()

    def migrate_rooms(self):
        eq = defaultdict(set)
        vc = defaultdict(set)
        for old_room_id, old_room in self.rb_root['Rooms'].iteritems():
            eq[old_room._locationName].update(e for e in old_room._equipment.split('`') if e)
            vc[old_room._locationName].update(e for e in getattr(old_room, 'avaibleVC', []) if e)

        print cformat('%{white!}migrating equipment')
        for name, eqs in eq.iteritems():
            l = Location.find_first(name=name)

            if l is None:
                print cformat('%{yellow!}*** WARNING')
                print cformat("%{{yellow!}}***%{{reset}} Location '{}' does not exist. Skipped equipment: {}".format(
                    name, eqs
                ))
                continue

            l.equipment_types.extend(EquipmentType(name=x) for x in eqs)
            print cformat('- [%{cyan}{}%{reset}] {}').format(name, eqs)
            db.session.add(l)
        db.session.commit()
        print

        print cformat('%{white!}migrating vc equipment')
        for name, vcs in vc.iteritems():
            l = Location.find_first(name=name)

            if l is None:
                print cformat('%{yellow!}*** WARNING')
                print cformat("%{{yellow!}}***%{{reset}} Location '{}' does not exist. Skipped VC equipment: {}".format(
                    name, vcs
                ))
                continue

            pvc = l.get_equipment_by_name('Video conference')
            for vc_name in vcs:
                req = EquipmentType(name=vc_name)
                req.parent = pvc
                l.equipment_types.append(req)
                print cformat('- [%{cyan}{}%{reset}] {}').format(name, req.name)
            db.session.add(l)
        db.session.commit()
        print

        print cformat('%{white!}migrating rooms')

        for old_room_id, old_room in self.rb_root['Rooms'].iteritems():
            l = Location.find_first(name=old_room._locationName)

            if l is None:
                print cformat('%{yellow!}*** WARNING')
                print cformat("%{{yellow!}}***%{{reset}} Location '{}' does not exist. Skipped room '{}'".format(
                    old_room._locationName, old_room.id
                ))
                continue

            r = Room(
                id=old_room_id,
                name=convert_to_unicode((old_room._name or '').strip() or generate_name(old_room)),
                site=convert_to_unicode(old_room.site),
                division=convert_to_unicode(old_room.division),
                building=convert_to_unicode(old_room.building),
                floor=convert_to_unicode(old_room.floor),
                number=convert_to_unicode(old_room.roomNr),

                notification_before_days=((old_room.resvStartNotificationBefore or None)
                                          if getattr(old_room, 'resvStartNotification', False)
                                          else None),
                notification_for_responsible=getattr(old_room, 'resvNotificationToResponsible', False),
                notification_for_assistance=getattr(old_room, 'resvNotificationAssistance', False),

                reservations_need_confirmation=old_room.resvsNeedConfirmation,

                telephone=convert_to_unicode(getattr(old_room, 'telephone', None)),
                key_location=convert_to_unicode(getattr(old_room, 'whereIsKey', None)),

                capacity=getattr(old_room, 'capacity', None),
                surface_area=getattr(old_room, 'surfaceArea', None),
                latitude=getattr(old_room, 'latitude', None),
                longitude=getattr(old_room, 'longitude', None),

                comments=convert_to_unicode(getattr(old_room, 'comments', None)),

                owner_id=self.merged_avatars.get(old_room.responsibleId, old_room.responsibleId),

                is_active=old_room.isActive,
                is_reservable=old_room.isReservable,
                max_advance_days=int(old_room.maxAdvanceDays) if getattr(old_room, 'maxAdvanceDays', None) else None
            )

            print cformat('- [%{cyan}{}%{reset}] %{grey!}{:4}%{reset}  %{green!}{}%{reset}').format(l.name, r.id,
                                                                                r.name)

            for old_bookable_time in getattr(old_room, '_dailyBookablePeriods', []):
                r.bookable_hours.append(
                    BookableHours(
                        start_time=old_bookable_time._startTime,
                        end_time=old_bookable_time._endTime
                    )
                )
                print cformat('  %{blue!}Bookable:%{reset} {}').format(r.bookable_hours[-1])

            for old_nonbookable_date in getattr(old_room, '_nonBookableDates', []):
                r.nonbookable_periods.append(
                    NonBookablePeriod(
                        start_dt=old_nonbookable_date._startDate,
                        end_dt=old_nonbookable_date._endDate
                    )
                )
                print cformat('  %{blue!}Nonbookable:%{reset} {}').format(r.nonbookable_periods[-1])

            if self.photo_path:
                try:
                    with open(os.path.join(self.photo_path, 'large_photos',
                                           get_canonical_name_of(old_room) + '.jpg'), 'rb') as f:
                        large_photo = f.read()
                except Exception:
                    large_photo = None

                try:
                    with open(os.path.join(self.photo_path, 'small_photos',
                                           get_canonical_name_of(old_room) + '.jpg'), 'rb') as f:
                        small_photo = f.read()
                except Exception:
                    small_photo = None

                if large_photo and small_photo:
                    r.photo = Photo(data=large_photo, thumbnail=small_photo)
                    print cformat('  %{blue!}Photos')

            new_eq = []
            for old_equipment in ifilter(None, old_room._equipment.split('`') + old_room.avaibleVC):
                room_eq = l.get_equipment_by_name(old_equipment)
                new_eq.append(room_eq)
                r.available_equipment.append(room_eq)
            if new_eq:
                print cformat('  %{blue!}Equipment:%{reset} {}').format(', '.join(sorted(x.name for x in new_eq)))

            for attr_name, value in getattr(old_room, 'customAtts', {}).iteritems():
                value = convert_to_unicode(value)
                if not value or ('Simba' in attr_name and value == u'Error: unknown mailing list'):
                    continue
                attr_name = attribute_map.get(attr_name, attr_name).replace(' ', '-').lower()
                ca = l.get_attribute_by_name(attr_name)
                if not ca:
                    print cformat('  %{blue!}Attribute:%{reset} {} %{red!}not found').format(attr_name)
                    continue
                attr = RoomAttributeAssociation()
                attr.value = value
                attr.attribute = ca
                r.attributes.append(attr)
                print cformat('  %{blue!}Attribute:%{reset} {} = {}').format(attr.attribute.title, attr.value)

            l.rooms.append(r)
            db.session.add(l)
            print
        db.session.commit()

    def migrate_reservations(self):
        print cformat('%{white!}migrating reservations')
        i = 1
        for rid, v in self.rb_root['Reservations'].iteritems():
            room = Room.get(v.room.id)
            if room is None:
                print cformat('  %{red!}skipping resv for dead room {0.room.id}: {0.id} ({0._utcCreatedDT})').format(v)
                continue

            repeat_frequency, repeat_interval = RepeatMapping.convert_legacy_repeatability(v.repeatability)
            booked_for_id = getattr(v, 'bookedForId', None)

            r = Reservation(
                id=v.id,
                created_dt=as_utc(v._utcCreatedDT),
                start_dt=utc_to_local(v._utcStartDT),
                end_dt=utc_to_local(v._utcEndDT),
                booked_for_id=self.merged_avatars.get(booked_for_id, booked_for_id) or None,
                booked_for_name=convert_to_unicode(v.bookedForName),
                contact_email=convert_to_unicode(v.contactEmail),
                contact_phone=convert_to_unicode(getattr(v, 'contactPhone', None)),
                created_by_id=self.merged_avatars.get(v.createdBy, v.createdBy) or None,
                is_cancelled=v.isCancelled,
                is_accepted=v.isConfirmed,
                is_rejected=v.isRejected,
                booking_reason=convert_to_unicode(v.reason),
                rejection_reason=convert_to_unicode(getattr(v, 'rejectionReason', None)),
                repeat_frequency=repeat_frequency,
                repeat_interval=repeat_interval,
                uses_vc=getattr(v, 'usesAVC', False),
                needs_vc_assistance=getattr(v, 'needsAVCSupport', False),
                needs_assistance=getattr(v, 'needsAssistance', False)
            )

            for eq_name in getattr(v, 'useVC', []):
                eq = room.location.get_equipment_by_name(eq_name)
                if eq:
                    r.used_equipment.append(eq)

            occurrence_rejection_reasons = {}
            if getattr(v, 'resvHistory', None):
                for h in reversed(v.resvHistory._entries):
                    ts = as_utc(parse_dt_string(h._timestamp))

                    if len(h._info) == 2:
                        possible_rejection_date, possible_rejection_reason = h._info
                        m = re.match(r'Booking occurrence of the (\d{1,2} \w{3} \d{4}) rejected',
                                     possible_rejection_reason)
                        if m:


# values.py
"""Functions for converting values of DICOM data elements to proper python types
"""
# Copyright (c) 2010-2012 Darcy Mason
# This file is part of pydicom, relased under an MIT license.
#    See the file license.txt included with this distribution, also
#    available at https://github.com/darcymason/pydicom

from struct import unpack, calcsize

from pydicom.compat import in_py2
from pydicom import compat
# Because DS can be based on float or decimal, import whole module, not DS
#    directly, so it can be changed in user code and be updated here also
import pydicom.valuerep
from pydicom.valuerep import MultiString, DA, DT, TM

from pydicom.config import logger, datetime_conversion

if not in_py2:
    from pydicom.valuerep import PersonName3 as PersonName
else:
    from pydicom.valuerep import PersonName  # NOQA

from pydicom.multival import MultiValue
import pydicom.uid
from pydicom.tag import Tag, TupleTag
from pydicom.filereader import read_sequence
from io import BytesIO
from pydicom.charset import default_encoding, text_VRs

def convert_tag(byte_string, is_little_endian, offset=0):
    if is_little_endian:
        struct_format = "<HH"
    else:
        struct_format = ">HH"
    return TupleTag(unpack(struct_format, byte_string[offset:offset + 4]))

def convert_ATvalue(byte_string, is_little_endian, struct_format=None):
    """Read and return AT (tag) data_element value(s)"""
    length = len(byte_string)
    if length == 4:
        return convert_tag(byte_string, is_little_endian)
    # length > 4
    if length % 4 != 0:
        logger.warn("Expected length to be multiple of 4 for VR 'AT', "
                    "got length %d", length)
    return MultiValue(Tag, [convert_tag(byte_string, is_little_endian, offset=x)
                            for x in range(0, length, 4)])

def convert_DA_string(byte_string, is_little_endian, struct_format=None):
    """Read and return a DA value"""
    if datetime_conversion:
        if not in_py2:
            byte_string = byte_string.decode(encoding)
        length = len(byte_string)
        if length != 8:
            logger.warn("Expected length to be 8, got length %d", length)
        return DA(byte_string)
    else:
        return convert_string(byte_string, is_little_endian, struct_format)

def convert_DS_string(byte_string, is_little_endian, struct_format=None):
    """Read and return a DS value or list of values"""
    if not in_py2:
        byte_string = byte_string.decode(default_encoding)
    # Below, go directly to DS class instance rather than factory DS,
    # but need to ensure last string doesn't have blank padding (use strip())
    return MultiString(byte_string.strip(), valtype=pydicom.valuerep.DSclass)

def convert_DT_string(byte_string, is_little_endian, struct_format=None):
    """Read and return a DT value"""
    if datetime_conversion:
        if not in_py2:
            byte_string = byte_string.decode(encoding)
        length = len(byte_string)
        if length < 14 or length > 26:
            logger.warn("Expected length between 14 and 26, got length %d", length)
        return DT(byte_string)
    else:
        return convert_string(byte_string, is_little_endian, struct_format)

def convert_IS_string(byte_string, is_little_endian, struct_format=None):
    """Read and return an IS value or list of values"""
    if not in_py2:
        byte_string = byte_string.decode(default_encoding)
    return MultiString(byte_string, valtype=pydicom.valuerep.IS)

def convert_numbers(byte_string, is_little_endian, struct_format):
    """Read a "value" of type struct_format from the dicom file. "Value" can be more than one number"""
    endianChar = '><'[is_little_endian]
    bytes_per_value = calcsize("=" + struct_format)  # "=" means use 'standard' size, needed on 64-bit systems.
    length = len(byte_string)
    if length % bytes_per_value != 0:
        logger.warn("Expected length to be even multiple of number size")
    format_string = "%c%u%c" % (endianChar, length // bytes_per_value, struct_format)
    value = unpack(format_string, byte_string)
    if len(value) == 1:
        return value[0]
    else:
        return list(value)  # convert from tuple to a list so can modify if need to

def convert_OBvalue(byte_string, is_little_endian, struct_format=None):
    """Return the raw bytes from reading an OB value"""
    return byte_string

def convert_OWvalue(byte_string, is_little_endian, struct_format=None):
    """Return the raw bytes from reading an OW value rep

    Note: pydicom does NOT do byte swapping, except in
    dataset.pixel_array function
    """
    return convert_OBvalue(byte_string, is_little_endian)  # for now, Maybe later will have own routine

def convert_PN(byte_string, is_little_endian, struct_format=None, encoding=None):
    """Read and return string(s) as PersonName instance(s)"""

    # XXX - We have to replicate MultiString functionality here because we can't decode
    # easily here since that is performed in PersonNameUnicode
    if byte_string and (byte_string.endswith(b' ') or byte_string.endswith(b'\x00')):
        byte_string = byte_string[:-1]

    splitup = byte_string.split(b"\\")

    if encoding and not in_py2:
        args = (encoding,)
    else:
        args = ()

    # We would like to return string literals
    if not in_py2:
        valtype = lambda x: PersonName(x, *args).decode()
    else:
        valtype = lambda x: PersonName(x, *args)

    if len(splitup) == 1:
        return valtype(splitup[0])
    else:
        return MultiValue(valtype, splitup)

def convert_string(byte_string, is_little_endian, struct_format=None, encoding=default_encoding):
    """Read and return a string or strings"""
    if not in_py2:
        byte_string = byte_string.decode(encoding)
    return MultiString(byte_string)

def convert_single_string(byte_string, is_little_endian, struct_format=None,
                          encoding=default_encoding):
    """Read and return a single string (backslash character does not split)"""
    if not in_py2:
        byte_string = byte_string.decode(encoding)
    if byte_string and byte_string.endswith(' '):
        byte_string = byte_string[:-1]
    return byte_string

def convert_SQ(byte_string, is_implicit_VR, is_little_endian,
               encoding=default_encoding, offset=0):
    """Convert a sequence that has been read as bytes but not yet parsed."""
    fp = BytesIO(byte_string)
    seq = read_sequence(fp, is_implicit_VR, is_little_endian, len(byte_string),
                        encoding, offset)
    return seq

def convert_TM_string(byte_string, is_little_endian, struct_format=None):
    """Read and return a TM value"""
    if datetime_conversion:
        if not in_py2:
            byte_string = byte_string.decode(encoding)
        length = len(byte_string)
        if length < 2 or length > 16:
            logger.warn("Expected length between 2 and 16, got length %d", length)
        return TM(byte_string)
    else:
        return convert_string(byte_string, is_little_endian, struct_format)

def convert_UI(byte_string, is_little_endian, struct_format=None):
    """Read and return a UI values or values"""
    # Strip off 0-byte padding for even length (if there)
    if not in_py2:
        byte_string = byte_string.decode(default_encoding)
    if byte_string and byte_string.endswith('\0'):
        byte_string = byte_string[:-1]
    return MultiString(byte_string, pydicom.uid.UID)

def convert_UN(byte_string, is_little_endian, struct_format=None):
    """Return a byte string for a VR of 'UN' (unknown)"""
    return byte_string

def convert_value(VR, raw_data_element, encoding=default_encoding):
    """Return the converted value (from raw bytes) for the given VR"""
    if VR not in converters:
        raise NotImplementedError("Unknown Value Representation '{0}'".format(VR))

    # Look up the function to convert that VR
    # Dispatch two cases: a plain converter, or a number one which needs a format string
    if isinstance(converters[VR], tuple):
        converter, num_format = converters[VR]
    else:
        converter = converters[VR]
        num_format = None

    # Ensure that encoding is in the proper 3-element format
    if isinstance(encoding, compat.string_types):
        encoding = [encoding, ] * 3

    byte_string = raw_data_element.value
    is_little_endian = raw_data_element.is_little_endian
    is_implicit_VR = raw_data_element.is_implicit_VR

    # Not only two cases. Also need extra info if is a raw sequence
    # Pass the encoding to the converter if it is a specific VR
    if VR == 'PN':
        value = converter(byte_string, is_little_endian, encoding=encoding)
    elif VR in text_VRs:
        # Text VRs use the 2nd specified encoding
        value = converter(byte_string, is_little_endian, encoding=encoding[1])
    elif VR != "SQ":
        value = converter(byte_string, is_little_endian, num_format)
    else:
        value = convert_SQ(byte_string, is_implicit_VR, is_little_endian,
                           encoding, raw_data_element.value_tell)
    return value

# converters map a VR to the function to read the value(s).
# for convert_numbers, the converter maps to a tuple (function, struct_format)
#                        (struct_format in python struct module style)
converters = {
    'UL': (convert_numbers, 'L'),


# Created By: Virgil Dupras
# Created On: 2009-04-25
# Copyright 2015 Hardcoded Software (http://www.hardcoded.net)
#
# This software is licensed under the "GPLv3" License as described in the "LICENSE" file,
# which should be included with this package. The terms are also available at
# http://www.gnu.org/licenses/gpl-3.0.html

from PyQt5.QtCore import Qt, QRect
from PyQt5.QtWidgets import (
    QMainWindow,
    QMenu,
    QLabel,
    QFileDialog,
    QMenuBar,
    QWidget,
    QVBoxLayout,
    QAbstractItemView,
    QStatusBar,
    QDialog,
    QPushButton,
    QCheckBox,
    QDesktopWidget,
)

from hscommon.trans import trget
from qtlib.util import moveToScreenCenter, horizontalWrap, createActions
from qtlib.search_edit import SearchEdit

from core.app import AppMode
from .results_model import ResultsView
from .stats_label import StatsLabel
from .prioritize_dialog import PrioritizeDialog
from .se.results_model import ResultsModel as ResultsModelStandard
from .me.results_model import ResultsModel as ResultsModelMusic
from .pe.results_model import ResultsModel as ResultsModelPicture

tr = trget("ui")

class ResultWindow(QMainWindow):
    def __init__(self, parent, app, **kwargs):
        super().__init__(parent, **kwargs)
        self.app = app
        self._setupUi()
        if app.model.app_mode == AppMode.Picture:
            MODEL_CLASS = ResultsModelPicture
        elif app.model.app_mode == AppMode.Music:
            MODEL_CLASS = ResultsModelMusic
        else:
            MODEL_CLASS = ResultsModelStandard
        self.resultsModel = MODEL_CLASS(self.app, self.resultsView)
        self.stats = StatsLabel(app.model.stats_label, self.statusLabel)
        self._update_column_actions_status()

        self.menuColumns.triggered.connect(self.columnToggled)
        self.resultsView.doubleClicked.connect(self.resultsDoubleClicked)
        self.resultsView.spacePressed.connect(self.resultsSpacePressed)
        self.detailsButton.clicked.connect(self.actionDetails.triggered)
        self.dupesOnlyCheckBox.stateChanged.connect(self.powerMarkerTriggered)
        self.deltaValuesCheckBox.stateChanged.connect(self.deltaTriggered)
        self.searchEdit.searchChanged.connect(self.searchChanged)
        self.app.willSavePrefs.connect(self.appWillSavePrefs)

    def _setupActions(self):
        # (name, shortcut, icon, desc, func)
        ACTIONS = [
            ("actionDetails", "Ctrl+I", "", tr("Details"), self.detailsTriggered),
            ("actionActions", "", "", tr("Actions"), self.actionsTriggered),
            (
                "actionPowerMarker",
                "Ctrl+1",
                "",
                tr("Show Dupes Only"),
                self.powerMarkerTriggered,
            ),
            ("actionDelta", "Ctrl+2", "", tr("Show Delta Values"), self.deltaTriggered),
            (
                "actionDeleteMarked",
                "Ctrl+D",
                "",
                tr("Send Marked to Recycle Bin..."),
                self.deleteTriggered,
            ),
            (
                "actionMoveMarked",
                "Ctrl+M",
                "",
                tr("Move Marked to..."),
                self.moveTriggered,
            ),
            (
                "actionCopyMarked",
                "Ctrl+Shift+M",
                "",
                tr("Copy Marked to..."),
                self.copyTriggered,
            ),
            (
                "actionRemoveMarked",
                "Ctrl+R",
                "",
                tr("Remove Marked from Results"),
                self.removeMarkedTriggered,
            ),
            (
                "actionReprioritize",
                "",
                "",
                tr("Re-Prioritize Results..."),
                self.reprioritizeTriggered,
            ),
            (
                "actionRemoveSelected",
                "Ctrl+Del",
                "",
                tr("Remove Selected from Results"),
                self.removeSelectedTriggered,
            ),
            (
                "actionIgnoreSelected",
                "Ctrl+Shift+Del",
                "",
                tr("Add Selected to Ignore List"),
                self.addToIgnoreListTriggered,
            ),
            (
                "actionMakeSelectedReference",
                "Ctrl+Space",
                "",
                tr("Make Selected into Reference"),
                self.app.model.make_selected_reference,
            ),
            (
                "actionOpenSelected",
                "Ctrl+O",
                "",
                tr("Open Selected with Default Application"),
                self.openTriggered,
            ),
            (
                "actionRevealSelected",
                "Ctrl+Shift+O",
                "",
                tr("Open Containing Folder of Selected"),
                self.revealTriggered,
            ),
            (
                "actionRenameSelected",
                "F2",
                "",
                tr("Rename Selected"),
                self.renameTriggered,
            ),
            ("actionMarkAll", "Ctrl+A", "", tr("Mark All"), self.markAllTriggered),
            (
                "actionMarkNone",
                "Ctrl+Shift+A",
                "",
                tr("Mark None"),
                self.markNoneTriggered,
            ),
            (
                "actionInvertMarking",
                "Ctrl+Alt+A",
                "",
                tr("Invert Marking"),
                self.markInvertTriggered,
            ),
            (
                "actionMarkSelected",
                "",
                "",
                tr("Mark Selected"),
                self.markSelectedTriggered,
            ),
            (
                "actionExportToHTML",
                "",
                "",
                tr("Export To HTML"),
                self.app.model.export_to_xhtml,
            ),
            (
                "actionExportToCSV",
                "",
                "",
                tr("Export To CSV"),
                self.app.model.export_to_csv,
            ),
            (
                "actionSaveResults",
                "Ctrl+S",
                "",
                tr("Save Results..."),
                self.saveResultsTriggered,
            ),
            (
                "actionInvokeCustomCommand",
                "Ctrl+Alt+I",
                "",
                tr("Invoke Custom Command"),
                self.app.invokeCustomCommand,
            ),
        ]
        createActions(ACTIONS, self)
        self.actionDelta.setCheckable(True)
        self.actionPowerMarker.setCheckable(True)

    def _setupMenu(self):
        self.menubar = QMenuBar()
        self.menubar.setGeometry(QRect(0, 0, 630, 22))
        self.menuFile = QMenu(self.menubar)
        self.menuFile.setTitle(tr("File"))