# -*- coding: utf-8 -*-
"""
copyright 2024 Peak Design
this file is part of hsr1, which is distributed under the GNU Lesser General Public License v3 (LGPL)
"""
import uuid
import os
import importlib_resources

import numpy as np
import pandas as pd

pvlib_imported = False

try:
    import pvlib
    pvlib_imported = True
except:
    print("pvlib was not successfully imported")

import hsr1
from hsr1.utils.spectrum import spectrumUtils

SCALING_FACTOR = 0.617

rand = np.random.default_rng()

class SyntheticDataset():
    def __init__(self):
        if not pvlib_imported:
            raise ImportError("pvlib is required to generate synthetic datasets")
    
    def __init__(self, spectral_func=None, 
                 start_date="2024-01-01 00:00:00+00:00", 
                 end_date="2024-01-02 00:00:00+00:00",
                 latitude=53.14,
                 longitude=-1.64,
                 altitude=40):
        
        
        if spectral_func is None:
            self.spectral_func = self.generate_spectral_data
        
        self.spectral_data = None
        self.system_data = None
        self.deployment_metadata = None
        self.accessory_data = None
        
        self.start_date = start_date
        self.end_date = end_date
        self.latitude = latitude
        self.longitude = longitude
        self.altitude = altitude
        
        self.dataseries_id = None
    
    def get_dfs(self):
        return [self.spectral_data, self.system_data, self.deployment_metadata, self.accessory_data]
    
    def generate_custom_column(self, table, column, function, *args, **kwargs):
        if table == "spectral_data":
            self.spectral_data[column] = function(self.spectral_data["pc_time_end_measurement"], 
                                                  *args, **kwargs)
        elif table == "system_data":
            self.system_data[column] = function(self.spectral_data["pc_time_end_measurement"], 
                                                *args, **kwargs)
        elif table == "accessory_data":
            self.accessory_data[column] = function(self.accessory_data["pc_time_end_measurement"], 
                                                   *args, **kwargs)
        
        elif table == "deployment_metadata":
            self.deployment_metadata[column] = function(*args, **kwargs)
            
        else:
            raise ValueError("table not found")
    
    def remove_rows(self, table, function, *args, **kwargs):
        if table == "spectral_data":
            self.spectral_data = self.spectral_data.loc[function(self.spectral_data["pc_time_end_measurement"], *args, **kwargs), :]
        if table == "system_data":
            self.system_data = self.system_data.loc[function(self.system_data["pc_time_end_measurement"], *args, **kwargs), :]
        if table == "accessory_data":
            self.accessory_data = self.accessory_data.loc[function(self.accessory_data["pc_time_end_measurement"], *args, **kwargs), :]

    def global_integral(self, timestamps):
        if self.system_data is None:
            raise ValueError("system_data is required, generate that first")
            
        
        
    
    def generate_spectral_data(self, period="min"):
        spectral_data = pd.DataFrame()
        
        spectral_data["pc_time_end_measurement"] = pd.date_range(self.start_date, self.end_date, freq=period).astype(str)
        num_readings = len(spectral_data)
        spectral_data["sample_id"] = self.generate_sample_ids(spectral_data["pc_time_end_measurement"])
        spectral_data["dataseries_id"] = self.generate_dataseries_id()
        spectral_data["global_spectrum"] = self.smarts_integral_data(spectral_data["pc_time_end_measurement"]).tobytes()
        spectral_data["diffuse_spectrum"] = self.smarts_integral_data(spectral_data["pc_time_end_measurement"]).tobytes()
        spectral_data["global_integral"] = pd.to_datetime(spectral_data["pc_time_end_measurement"]).dt.hour
        spectral_data["diffuse_integral"] = spectral_data["global_integral"]/2
        spectral_data["global_molar"] = spectral_data["global_integral"]
        spectral_data["diffuse_molar"] = spectral_data["diffuse_integral"]
        spectral_data["camera_temp"] = rand.normal(30, 0.1, num_readings)
        
        
        self.spectral_data = spectral_data
    
    def generate_system_data(self, period="min"):
        system_data = pd.DataFrame()
        
        system_data["pc_time_end_measurement"] = pd.date_range(self.start_date, self.end_date, freq=period).astype(str)
        system_data["gps_time"] = pd.to_datetime(system_data["pc_time_end_measurement"]).dt.tz_localize(None).astype(str)
        system_data["gps_status"] = "A"
        system_data["gps_latitude"] = self.latitude
        system_data["gps_longitude"] = self.longitude
        system_data["gps_altitude"] = 100
        system_data["pressure"] = 1000
        system_data["baro_temp"] = 25
        system_data["rh"] = 13
        system_data["rh_temp"] = 0.0
        system_data["dataseries_id"] = self.generate_dataseries_id()
        system_data["sample_id"] = self.generate_sample_ids(system_data["pc_time_end_measurement"])
        
        
        self.system_data = system_data
    
    
    def generate_system_data_accessory(self, accessory_data=None, spectral_timestamps=None, spectral_sample_ids=None):
        if accessory_data is None:
            accessory_data = self.accessory_data
        if spectral_timestamps is None:
            spectral_timestamps = self.spectral_data["pc_time_end_measurement"]
        if spectral_sample_ids is None:
            spectral_sample_ids = self.spectral_data["sample_id"]
        
        system_data = hsr1.utils.ReformatData().reformat_system_data_reformatted_accessory(accessory_data, spectral_timestamps, spectral_sample_ids)
        
        system_data["dataseries_id"] = self.generate_dataseries_id()
        system_data["sample_id"] = self.generate_sample_ids(system_data["pc_time_end_measurement"])
        
        self.system_data = system_data
        
        
    
    def generate_deployment_metadata(self, filepath):
        self.deployment_metadata = hsr1.utils.ReformatData().reformat_deployment_metadata(filepath)
        # self.deployment_metadata["default_elevation"] = str(self.altitude)
        # self.deployment_metadata["default_latitude"] = str(self.latitude)
        # self.deployment_metadata["default_longitude"] = str(self.longitude)
    
    def generate_accessory_data(self, period="min"):
        accessory_data = pd.DataFrame()
        
        accessory_data["pc_time_end_measurement"] = pd.date_range(self.start_date, self.end_date, freq=period).astype(str)
        
        num_readings = len(accessory_data)
        
        accessory_data["Clock_Error"] = 0.0
        accessory_data["Latitude"] = rand.normal(self.latitude, 0.1, num_readings)
        accessory_data["Longitude"] = rand.normal(self.longitude, 0.1, num_readings)
        accessory_data["Altitude"] = rand.normal(100.0, 0.1, num_readings)
        accessory_data["GPSAge"] = 0.0
        accessory_data["NSV"] = rand.normal(10.0, 0.1, num_readings).astype(int)
        accessory_data["GPSSpare"] = rand.normal(5.0, 0.1, num_readings)
        accessory_data["_15Vin"] = rand.normal(23000.0, 0.1, num_readings)
        accessory_data["_3VCAP"] = rand.normal(22000.0, 0.1, num_readings)
        accessory_data["HTRIn"] = rand.normal(50.0, 0.1, num_readings)
        accessory_data["Vcc"] = rand.normal(3300.0, 0.1, num_readings)
        accessory_data["VSpare"] = rand.normal(0.0, 0.1, num_readings)
        accessory_data["I_Tot"] = rand.normal(484.0, 0.1, num_readings)
        accessory_data["I_15VPC"] = rand.normal(200.0, 0.1, num_readings)
        accessory_data["I_15VCAM"] = rand.normal(100.0, 0.1, num_readings)
        accessory_data["I_5VPC"] = rand.normal(4.0, 0.1, num_readings)
        accessory_data["ISpare"] = 0.0
        accessory_data["T_CPU"] = rand.normal(30.0, 0.1, num_readings)
        accessory_data["T_Bezel"] = rand.normal(29.0, 0.1, num_readings)
        accessory_data["T_RH"] = 0.0
        accessory_data["T_Baro"] = rand.normal(27.0, 0.1, num_readings)
        accessory_data["RH"] = rand.normal(12.0, 0.1, num_readings)
        accessory_data["Pressure"] = rand.normal(900.0, 0.1, num_readings)
        accessory_data["IMU_Output_Type"] = 2.0
        accessory_data["Yaw"] = rand.normal(358.0, 0.1, num_readings)
        accessory_data["Roll"] = rand.normal(-1.62, 0.1, num_readings)
        accessory_data["Pitch"] = rand.normal(10.31, 0.1, num_readings)
        accessory_data["Q_W"] = 16314.0
        accessory_data["Q_X"] = -1476.0
        accessory_data["Q_Y"] = 219.0
        accessory_data["Q_Z"] = 217.0
        accessory_data["Acc_X"] = 0.02
        accessory_data["Acc_Y"] = -0.23
        accessory_data["Acc_Z"] = 0.06
        accessory_data["IMU1"] = 0.0
        accessory_data["IMU2"] = 0.0
        accessory_data["Control_flags"] = 130.0
        accessory_data["Control_flags2"] = 0.0
        accessory_data["StatusFlags"] = 6144.0
        accessory_data["Heater"] = 0.0
        accessory_data["MsgCount"] = 69.0
        accessory_data["RdgCount"] = 0.0
        accessory_data["Watchdog"] = 3590.0
        accessory_data["gps_time"] = pd.to_datetime(accessory_data["pc_time_end_measurement"]).dt.tz_localize(None).astype(str)
        accessory_data["cpu_time"] = pd.to_datetime(accessory_data["pc_time_end_measurement"]).dt.tz_localize(None).astype(str)
        accessory_data["dataseries_id"] = self.generate_dataseries_id()
        accessory_data["sample_id"] =  self.generate_sample_ids(accessory_data["pc_time_end_measurement"])
        
        self.accessory_data = accessory_data
    
    
    """
    ---------------------------------------------------------------------------
    individual column functions
    ---------------------------------------------------------------------------
    """
    
    
    def generate_sample_ids(self, timeseries):
        """must have spectral_data generated first"""
        new_sample_ids = pd.DataFrame()
        new_sample_ids["pc_time_end_measurement"] = timeseries
        
        
        if self.spectral_data is not None:
            new_sample_ids = new_sample_ids.merge(self.spectral_data[["pc_time_end_measurement", "sample_id"]], how="left", on="pc_time_end_measurement")
        
        if "sample_id" in new_sample_ids.columns:
            not_matched = pd.isnull(new_sample_ids["sample_id"])
            new_sample_ids.loc[not_matched, "sample_id"] = [str(uuid.uuid1()) for i in range(np.sum(not_matched))]
        else:
            new_sample_ids["sample_id"] = [str(uuid.uuid1()) for i in range(len(timeseries))]
        
        return new_sample_ids["sample_id"]
    
    def generate_dataseries_id(self):
        if self.dataseries_id is None:
            self.dataseries_id = str(uuid.uuid1())
        
        return self.dataseries_id
    
    def linear_integral(self, timeseries):
        return np.ones(len(timeseries))

    def dist_from_midday(self, timeseries):
        timeseries = pd.to_datetime(timeseries)
        mins = timeseries.dt.hour*60 + timeseries.dt.minute
        return 720-np.abs(mins-720)

    def pvlib_integral_static(self, timeseries, *args):
        location = pvlib.location.Location(self.latitude, 
                                           self.longitude, 
                                           altitude=self.altitude)
        integrals = location.get_clearsky(pd.DatetimeIndex(timeseries))
        data = integrals[["ghi", "dhi"]].values*SCALING_FACTOR
        
        return data
    
    def smarts_integral_data(self, timeseries, filepath=None):
        if filepath is None:
            res = importlib_resources.files("hsr1.data").joinpath("smarts.txt")
            file = importlib_resources.as_file(res)
            with file as f:
                filepath = f
        
        smarts_data = spectrumUtils.read_simple_file(filepath)
        
        return smarts_data*len(timeseries)
    
    def integral_hour_of_day(self, timeseries):
        timeseries = pd.to_datetime(timeseries)
        return np.array([timeseries.dt.hour*100, timeseries.dt.hour*50]).T
    
    def remove_n_hour(self, timeseries, hour):
        timestamps = pd.to_datetime(timeseries)
        return timestamps.dt.hour != hour
    
    def num_readings(self, timeseries, base=0, scale=1):
        data = [i for i in range(len(timeseries))]
        data = base + np.array(data)*scale
        return data
        
        
"""
GUIDE TO WRITING CUSTOM FUNCTIONS:

dataseries and timeseries is automatically passed in, so the function must expect it
return a list, np array or series of data for the specific column

use it with

dataseries.generate_custom_column(table, column, function, **kwargs)
"""





if __name__ == "__main__":
    db_name = "databases/synthetic.db"
    
    if os.path.exists(db_name):
        os.remove(db_name)
    
    latitude = 13.19
    longitude = -59.54
    altitude = 40
    
    synthetic_dataset = SyntheticDataset(start_date="2023-01-01 00:00:00+00:00", 
                                         end_date="2023-01-03 23:59:00+00:00",
                                         latitude=latitude, longitude=longitude,
                                         altitude=40)
    
    synthetic_dataset.generate_spectral_data("10s")
    synthetic_dataset.generate_system_data()
    synthetic_dataset.generate_deployment_metadata("C:/Users/albie/work/Datasets/Winster 2023/HSR1-002 Winster 2024 Deployment.ini")
    synthetic_dataset.generate_accessory_data()
    
    synthetic_dataset.generate_custom_column("spectral_data", 
                                             ["global_integral", "diffuse_integral"], 
                                             synthetic_dataset.pvlib_integral_static)
    
    dfs = synthetic_dataset.get_dfs()
    
    db = hsr1.DBDriver(db_name)
    db.store(dfs)
    graph = hsr1.Graph(db)
    graph.plot_integral()
    graph.plot_accessory()
    graph.plot_gps()
    graph.plot_dips_summary()