import numpy as np
import dataclasses
import base64
import struct
import xml.etree.ElementTree as ET
from typing import List
import sys

@dataclasses.dataclass
class PixCfg:
    mask_bit: int
    test_bit: int
    thl_adj: int

class quadpixXml:
    def __init__(self):
        self.data = None
        self.items = {}

    def list_items(self, theitem, items):
        # Recursive function to traverse XML elements and populate the items dictionary
        for item in theitem:
            if len(item):
                newitems = {}
                self.list_items(item, newitems)
                items[item.tag] = newitems
            elif not item.text:
                items[item.tag] = ""
            elif len(item.text) > 1000:
                # Decode base64 data and unpack if needed
                data = base64.b64decode(item.text)
                if "calib" in item.tag:
                    data = struct.unpack('d' * (len(data) // 8), data)
                else:
                    data = struct.unpack('B' * len(data), data)
                items[item.tag] = data
            else:
                items[item.tag] = item.text

    def load(self, file_name: str):
        # Parse and load XML file into the items dictionary
        tree = ET.parse(file_name)
        root = tree.getroot()
        self.list_items(root, self.items)

    def get_data(self, section: str, name: str):
        # Retrieve data from the items dictionary
        return self.items[section][name]

    def get_binary_pix_cfg(self, chip_id: str) -> List[PixCfg]:
        # Extract and convert BinaryPixelCfg data into a list of PixCfg objects
        data = self.get_data(chip_id, "BinaryPixelCfg")
        result = []
        for x in data:
            result.append(PixCfg(mask_bit=x & 0x1, test_bit=(x >> 5) & 0x1, thl_adj=(x >> 1) & 0xF))
        return result

    def tpx3_restore_chip_config(self, filename, chip_id, chip_array=False):
        # Load the XML file
        self.load(filename)

        # Finding all settings   
        b_layout = self.get_layouts()

        # Finding all settings   
        b_settings = self.get_tpx3_settings() 

        system_dacs = []
        system_info = []

        if chip_array:
            # If restoring multiple chips
            for id in chip_id:
                system_dacs.append(self.get_chip_dacs(id))
                system_info.append(self.get_chip_info(id))
        else:
            # If restoring a single chip
            system_dacs.append(self.get_chip_dacs(chip_id))
            system_info.append(self.get_chip_info(chip_id))
           
        return b_layout, b_settings, system_dacs, system_info

    def get_layouts(self):
        # Extract layout information from the items dictionary
        tpx3_layout = {
            'Angles': self.items['Layout']['Angles'], 
            'Chips': self.items['Layout']['Chips'], 
            'Width': self.items['Layout']['Width'], 
            'Height': self.items['Layout']['Height']
        }
        return tpx3_layout

    def get_tpx3_settings(self):
        # Extract TPX3 settings from the items dictionary
        tpx3_settings = {
            'Bias': self.items['Settings']['Bias'], 
            'SensorRefreshEnabled': self.items['Settings']['SensorRefreshEnabled'], 
            'SensorRefreshTime': self.items['Settings']['SensorRefreshTime'], 
            'ExtBiasSerial': self.items['Settings']['ExtBiasSerial'], 
            'ExtBiasSrcIndex': self.items['Settings']['ExtBiasSrcIndex'], 
            'InterpolateMaskedPixels': self.items['Settings']['InterpolateMaskedPixels'], 
            'InterpolateMaskedPixelsFlags': self.items['Settings']['InterpolateMaskedPixelsFlags'], 
            'ReleativeTHL': self.items['Settings']['ReleativeTHL'], 
            'UseCalibration': self.items['Settings']['UseCalibration'], 
            'Polarity': self.items['Settings']['Polarity'], 
            'ConvertToaTime': self.items['Settings']['ConvertToaTime'], 
            'OperationMode': self.items['Settings']['OperationMode']
        }
        return tpx3_settings
   

    def get_chip_dacs(self, chip_id):
        # Extract DACs information for a chip from the items dictionary
        chip_dacs = np.array([
            int(self.items[chip_id]['PreampOn']),
            int(self.items[chip_id]['PreampOff']),
            int(self.items[chip_id]['NCas']),
            int(self.items[chip_id]['Ikrum']),
            int(self.items[chip_id]['Fbk']),
            int(self.items[chip_id]['Threshold']),
            int(self.items[chip_id]['THLFine']),
            int(self.items[chip_id]['THLCoarse']),
            int(self.items[chip_id]['DiscS1On']),
            int(self.items[chip_id]['DiscS1Off']),
            int(self.items[chip_id]['DiscS2On']),
            int(self.items[chip_id]['DiscS2Off']),
            int(self.items[chip_id]['PixelDac']),
            int(self.items[chip_id]['TpBuffIn']),
            int(self.items[chip_id]['TpBuffOut']),
            int(self.items[chip_id]['TpCoarse']),
            int(self.items[chip_id]['TpFine']),
            int(self.items[chip_id]['CpPll']),
            int(self.items[chip_id]['PllVCntrl']),
            int(self.items[chip_id]['ThresholdCalibCoeffA']),
            int(self.items[chip_id]['ThresholdCalibCoeffB']),
            int(self.items[chip_id]['MinThreshold'])
        ])
        return chip_dacs


    def get_chip_info(self, chip_id):
        # Extract chip information from the items dictionary
        chip_info = [
            self.items[chip_id]['Polarity'],
            self.items[chip_id]['SensorType'],
            int(self.items[chip_id]['SensorThickness']),
            int(self.items[chip_id]['SensorPitch'])
        ]
        return chip_info


def main():
    # Check if both filename and zchip_ids are provided as command line arguments
    if len(sys.argv) < 3:
        print("Usage: python script.py <filename> <zchip_ids>")
        sys.exit(1)

    # Get filename and zchip_ids from command line arguments
    filename = sys.argv[1]
    zchip_ids = sys.argv[2].split(',')

    yy = quadpixXml()
    yy.load(filename)

    # Extract binary pixel configurations and print the first 10 elements
    pixcfg = yy.get_binary_pix_cfg(zchip_ids[0])
    print("First 10 PixCfg:")
    for idx in range(10):
        print(f'Mask = {pixcfg[idx].mask_bit}, Test Bit = {pixcfg[idx].test_bit}, THL Adj = {pixcfg[idx].thl_adj}')

    # Restore chip configuration
    chip_id_to_restore = zchip_ids[0]  # Replace with the actual chip_id you want to restore
    b_layout, b_settings, system_dacs, system_info = yy.tpx3_restore_chip_config(filename, chip_id_to_restore)

    # Print restored DACs
    print('')
    print('Restored DACs = ', system_dacs)

if __name__ == "__main__":
    main()