#!/usr/bin/env python2 
# coding: utf-8
"""
authors: Max Ahnen, Dominik Neise
----------------------------------------------------------------------------
"THE BEER-WARE LICENSE" (Revision 42):
Max Ahnen and Dominik Neise wrote this file. As long as you retain this notice you
can do whatever you want with this stuff. If we meet some day, and you think
this stuff is worth it, you can buy us a beer.
----------------------------------------------------------------------------

This script calculates the mean of the SQM magnitude for a given run
(run is given by its start_ and stop_time)
and the p-value for a linear fit.

We hope that this will be used to separate good data, where 
the linear fit is perfect (p-value > 1e-3?) and bad fits, 
where clouds let the magnitude brightness fluctuate stronger,
so the fit worsens.

Returns: a string
  "results {mean magnitude:f} {p-value:f}"   
"""

from astropy.io import fits
import numpy as np
import scipy as sp
import sys
import ROOT 

import pandas as pd
from sqlalchemy import create_engine

import glob

from calendar import timegm
import time

database = {
    'user': 'factread',
    'password': 'r3adfac!',
    'host': '129.194.168.95',
    'table': 'factdata',
}

root_database = {
    'user': 'root',
    'password': '1440Gapd',
    'host': 'localhost',
    'table': 'factdata',
}


db_string = '{user}:{password}@{host}/{table}'

try:
    factdb = create_engine('mysql+mysqldb://'+ db_string.format(**database))
    factdb_root = create_engine('mysql+mysqldb://'+ db_string.format(**root_database))
except ImportError:
    factdb = create_engine('mysql+pymysql://'+db_string.format(**database))
    factdb_root = create_engine('mysql+pymysql://'+db_string.format(**root_database))


def get_list_of_SQM_files(base_path='/daq/aux'):
    return sorted(glob.glob(base_path+'/*/*/*/*.SQM_CONTROL_DATA.fits'))

def get_y_m_d(file_path):
    s = file_path.split('/')[-1].split('.')[0]
    return s[0:4], s[4:6], s[6:8]

def get_night(file_path):
    return file_path.split('/')[-1].split('.')[0]

def mag_mean_p_value(fits_file_path, start_time, stop_time):

	d = fits.open(fits_file_path)[1].data
	d = d[(d['Time'] > start_time) * (d['Time'] < stop_time)]
        if len(d)==0:
            return None

	x =  d['Time'].copy() - d['Time'][0]
        y = d['Mag'].astype(np.float64).copy()
        start_time, stop_time = x[0], x[-1]

        sigma  = 0.025 / np.sqrt(12.)
	g = ROOT.TGraphErrors(len(x), x, y, np.zeros(len(x), dtype=np.float64), np.ones(len(x), dtype=np.float64)*sigma)
	f = ROOT.TF1("linfit", "pol1", start_time, stop_time)
	g.Fit(f, "E")
        
        function = f.Eval
	y_fit = map(function,  x)
 
        result = {
            'fSqmMagMean' : d['Mag'].mean(),
            'fSqmMagLinFitPValue' : f.GetProb(),
            'fSqmMagLinFitChi2' : f.GetChisquare(),
            'fSqmMagLinFitNdf' : f.GetNDF(),
            'fSqmMagLinFitSlope' : f.GetParameters()[1],
        }

	return result

def timestamp_from_string(stri):
    return timegm(time.strptime(stri.replace('Z', 'UTC'),'%Y-%m-%d %H:%M:%S%Z'))

def update_dict_in_database(some_dict, db, primary_keys=("fNight", "fRunID")):
    commands = []
    commands.append('BEGIN;')

    # "UPDATE RunInfo SET weight = 160, desiredWeight = 145 WHERE id = 1;"
    update_string = "UPDATE RunInfo SET "
    first = True
    for k in some_dict:
        if k not in primary_keys:
            if not first:
                update_string += ', '
            else:
                first = False

            update_string += "{0} = {1}".format(k, some_dict[k])
    update_string += " WHERE "
    first = True
    for k in some_dict:
        if k in primary_keys:
            if not first:
                update_string += ' AND '
            else:
                first = False
            update_string += "{0} = {1}".format(k, some_dict[k])
    update_string += ';'

    commands.append(update_string)
    commands.append('COMMIT;')

    # print commands
    for com in commands:
        db.engine.execute(com)

    
if __name__ == "__main__":
    aux_files = get_list_of_SQM_files()
    for aux_file in aux_files:
      night = get_night(aux_file)

      query = ("SELECT fRunID, fRunStart, fRunStop from RunInfo"
               " WHERE fNight={0}").format(night)

      df = pd.read_sql_query(query, factdb)

      for i in range(df.shape[0]):
          row = df.iloc[i]

          try:
              run_start = timestamp_from_string(str(row['fRunStart'])+'Z')/(24.*3600.)
              run_stop = timestamp_from_string(str(row['fRunStop'])+'Z')/(24.*3600.)
              run_id = row['fRunID']
          except ValueError as e:
              print e
              print row['fRunStart'], row['fRunStop']
              continue

          result= mag_mean_p_value(aux_file, run_start, run_stop)
          if result is None:
              continue
          result['fNight'] = int(night)
          result['fRunID'] = run_id

          update_dict_in_database(result, factdb_root)

          print time.asctime(), aux_file, run_start, run_stop, result