"""Add electron lifetime
"""
import datetime
from collections import defaultdict

import numpy as np
import os
import sympy
import pytz
import requests
from hax import slow_control
from pax import configuration, units
from sympy.parsing.sympy_parser import parse_expr

from cax import config
from cax.task import Task

PAX_CONFIG = configuration.load_configuration('XENON1T')


class CorrectionBase(Task):
    "Derive correction"

    collection_name = 'purity'

    def __init__(self):
        self.correction_collection = config.mongo_collection(self.collection_name)
        Task.__init__(self)

    def evaluate(self):
        raise NotImplementedError()

    def each_run(self):
        if self.key == 'slow_control' and 'slow_control' in self.run_doc:
            return
        else:
            short_key = self.key.split('.')[-1]
            if 'processor' in self.run_doc and \
                  'DEFAULT' in self.run_doc['processor'] and \
                   short_key in self.run_doc['processor']['DEFAULT']:
                return

        if 'end' not in self.run_doc:
            return

        self.get_correction()

        if not config.DATABASE_LOG:
            return

        # Update run database
        self.collection.find_and_modify({'_id'   : self.run_doc['_id'],
                                         self.key: {'$exists': False}},
                                        {'$set': {self.key: self.evaluate()}})

    def get_correction(self):
        # Fetch the latest electron lifetime fit
        doc = self.correction_collection.find_one(sort=(('calculation_time',
                                                         -1),))
        # Get fit function
        self.function = parse_expr(doc['function'])

    def get_time_range(self):
        if self.run_doc['end'] < datetime.datetime(2016, 7, 20):
            dt = datetime.timedelta(minutes=30)
        else:
            dt = datetime.timedelta(minutes=3)

        return (self.run_doc['start'] - dt,
                self.run_doc['end'] + dt)


class AddElectronLifetime(CorrectionBase):
    """Copy data to here

    If data exists at a reachable host but not here, pull it.
    """
    key = 'processor.DEFAULT.electron_lifetime_liquid'
    correction_units = units.us

    def evaluate(self):
        # Compute lifetime from this function on this dataset
        subs = {"t": self.run_doc['start'].timestamp()}
        lifetime = self.function.evalf(subs=subs)
        lifetime = float(lifetime)  # Convert away from Sympy type.

        run_number = self.run_doc['number']
        self.log.info("Run %d: calculated lifetime of %d us" % (run_number,
                                                                lifetime))
        return lifetime * self.correction_units


class AddSlowControlInformation(CorrectionBase):
    """Add all slow control highlight information to run db
    """
    key = 'slow_control'
    correction_units = 1

    def get_correction(self):
        pass

    def evaluate(self):
        time_range = self.get_time_range()

        data = defaultdict(dict)
        for key1, value1 in slow_control.VARIABLES.items():
            for key2, value2 in value1.items():
                series = slow_control.get_series(value2, time_range)
                if len(series):
                    data[key1][key2] = float(series.iloc[0])
                else:
                    data[key1][key2] = np.nan

        return data


class AddGains(CorrectionBase):
    """Copy data to here

    If data exists at a reachable host but not here, pull it.
    """
    collection_name = 'gains'
    key = 'processor.DEFAULT.gains'
    correction_units = units.V  # should be 1

    def each_run(self):
        """Only run on data manager at LNGS
        """
        if config.get_hostname() == 'xe1t-datamanager':
            CorrectionBase.each_run(self)

    def evaluate(self):
        """Make an array of all PMT gains."""
        start = self.run_doc['start']
        timestamp = start.replace(tzinfo=pytz.utc).timestamp()

        self.log.info("Run %d: gains computing" % self.run_doc['number'])
        gains = self.get_gains(timestamp)

        self.log.info("Run %d: gains:" % self.run_doc['number'])
        self.log.info(gains)
        return gains

    def get_gains(self, timestamp):
        """Timestamp is a UNIX timestamp in UTC
        """
        V = sympy.symbols('V')
        pmt = sympy.symbols('pmt', integer=True)

        voltages = self.get_voltages(timestamp)
        gains = []
        for i, voltage in enumerate(voltages):
            self.log.debug("Deriving HV for PMT %d" % i)
            gain = self.function.evalf(subs={V  : voltage,
                                              pmt: i})
            gains.append(float(gain) * self.correction_units)

        return gains

    def get_voltages(self, time):
        try:
            r = requests.post('https://172.16.2.105:4040/WebService.asmx/getLastMeasuredPMTValues',
                          data = {'EndDateUnix' : int(time),
                                  'username':'slowcontrolwebserver',
                                  'api_key' : os.environ.get('api_key'),
                              },
                          verify=False)

            r.raise_for_status()
        except requests.exceptions.HTTPError as e:
            self.log.exception(e)
            self.log.info("Sleeping 10 seconds, then retrying")
            time.sleep(10)
            return self.get_voltages(time)

        self.log.info(time)
        self.log.info(r.url)
        self.log.info(r.text)
        pmts = slow_control.VARIABLES['pmts']
        mapping = {v: int(k.split('_')[1]) for k,v in pmts.items()}

        voltages = len(PAX_CONFIG['DEFAULT']['pmts'])*[0]

        for doc in r.json():
            if doc['tagname'] in mapping.keys():
                voltages[mapping[doc['tagname']]] = doc['value']

        return voltages


@@ -4435,24 +4435,56 @@
 eger=True)%0A%0A
+        # Grab voltages from SC%0A
         volt
@@ -4519,16 +4519,99 @@
 estamp)%0A
+%0A        if None in voltages:%0A            raise ValueError(%22Missing SC variable%22)%0A%0A
         
@@ -5858,17 +5858,20 @@
 mts'%5D)*%5B
-0
+None
 %5D%0A%0A     

Throw error if missing SC variable