# This file is part of Fast Servo Software Package.
#
# Copyright (C) 2023 Jakub Matyas
# Warsaw University of Technology <jakubk.m@gmail.com>
# SPDX-License-Identifier: GPL-3.0-or-later
#
# This program 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.
# 
# This program 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 this program. If not, see <https://www.gnu.org/licenses/>.

import time
import spidev
from pyfastservo.common import (
    ADC_AFE_CTRL_ADDR,
    ADC_BITSLIP_ADDR,
    ADC_CH0_HIGH_ADDR,
    ADC_CH0_LOW_ADDR,
    ADC_CH1_HIGH_ADDR,
    ADC_CH1_LOW_ADDR,
    ADC_DELAY_ADDR,
    ADC_FRAME_ADDR,
    AUX_ADC_ADDR,
    MAP_MASK,
    PAGESIZE,
    write_to_memory,
    read_from_memory
)

# /dev/spidev1.0  <=> spidev<BUS>.<DEVICE>
MAIN_ADC_BUS = 1
MAIN_ADC_DEVICE = 1

AUX_ADC_BUS = 1
AUX_ADC_PORT_A = 2
AUX_ADC_PORT_B = 3


def spi_write(spi, address, value):
    spi.xfer2([address, value])

def spi_read(spi, address):
    rx_buffer = spi.xfer2([0x80 | address, 0x00])
    return rx_buffer[1]

def main_adc_config(spi, test_pattern):
    high_word = (test_pattern & 0xFF00) >> 8
    low_word = test_pattern & 0xFF

    spi_write(spi, 0x00, 0x80)              # reset
    spi_write(spi, 0x01, 0x20)              # REGISTER A1: set to Two's complement Data Format
    spi_write(spi, 0x02, 0x15)              # REGISTER A2: set to LVDS output, set 4 data lanes and turn on test mode
    spi_write(spi, 0x03, high_word)         # REGISTER A3: test pattern high word
    spi_write(spi, 0x04, low_word)          # REGISTER A4: test pattern low word

def main_adc_test_mode(spi, enable):
    reg_contents = 0x15 if enable else 0x11 # set to LVDS output, set 4 data lanes and turn on or off test mode
    spi_write(spi, 0x02, reg_contents)

def verify_adc_registers(spi, reg_to_check):
    for register, expected_value in reg_to_check.items():
        value = spi_read(spi, register)
        print(f"Spi readback register 0x{register:02x}: 0x{value:02x}")
        if value != expected_value:
            print(f"Different value read than sent in reg 0x{register:02x}")

def read_frame():
    return read_from_memory(ADC_FRAME_ADDR, 1)[0]

def perform_bitslip():
    for i in range(4):
        current_frame = read_frame()
        if current_frame & 0x0F != 0x0C:
            print(f"Performing bitslip (iteration: {i}). Current frame: 0x{current_frame:02x}")
            write_to_memory(ADC_BITSLIP_ADDR, 1)
        else:
            print(f"No bitslip required; Current frame: 0x{current_frame:02x}")
            return

def mmcm_rst():
    curr_cfg = read_from_memory(ADC_AFE_CTRL_ADDR, 1)[0] & 0x0F
    write_to_memory(ADC_AFE_CTRL_ADDR, 0x10 | curr_cfg) # Reset MMCM
    write_to_memory(ADC_AFE_CTRL_ADDR, 0x00 | curr_cfg) # Release MMCM Reset
    while not(read_frame() & 0x10):
        print(f"Waiting for MMCM to lock")
        time.sleep(0.001)
    
def find_edge():
    prev_frame = read_frame()
    for tap_delay in range(32):
        write_to_memory(ADC_DELAY_ADDR, tap_delay)
        current_frame = read_frame()

        print(f"Tap delay: {tap_delay}, Current frame: 0x{current_frame:02x}")
        print(f"prev_frame: 0x{prev_frame:02x}")

        if current_frame != prev_frame:
            final_delay = ((tap_delay+1) // 2) + 2
            print(f"Edge detected; setting iDelay to: {final_delay}")
            write_to_memory(ADC_DELAY_ADDR, final_delay)
            return

        prev_frame = current_frame

    # If no edge detected
    final_delay = 11
    print(f"No edge detected; setting iDelay to: {final_delay}")
    write_to_memory(ADC_DELAY_ADDR, final_delay)

def read_adc_channel(high_addr, low_addr):
    return (read_from_memory(high_addr, 1)[0] << 8) | read_from_memory(low_addr, 1)[0]

def print_adc_channels():
    adc_ch0 = read_adc_channel(ADC_CH0_HIGH_ADDR, ADC_CH0_LOW_ADDR)
    adc_ch1 = read_adc_channel(ADC_CH1_HIGH_ADDR, ADC_CH1_LOW_ADDR)
    print(f"Final ADC_CH0: 0x{adc_ch0:04x}")
    print(f"Final ADC_CH1: 0x{adc_ch1:04x}")

def enable_adc_afe(ch1_x10=False, ch2_x10=False):
    ctrl_value = (ch2_x10 << 1) | ch1_x10
    write_to_memory(ADC_AFE_CTRL_ADDR, ctrl_value)
    afe_ctrl = read_from_memory(ADC_AFE_CTRL_ADDR, 1)[0]
    print(f"ADC_AFE_CTRL: 0x{afe_ctrl:02X}")
    return afe_ctrl

def configure_ltc2195():
    spi = spidev.SpiDev()
    try:
        spi.open(MAIN_ADC_BUS, MAIN_ADC_DEVICE)
        spi.max_speed_hz = 50000
        spi.mode = 0b00  # CPOL = 0 CPHA = 0
        spi.cshigh = False

        test_pattern = 0x811F
        main_adc_config(spi, test_pattern)

        verify_adc_registers(spi, {
            0x01: 0x20,
            0x02: 0x15,
            0x03: (test_pattern & 0xFF00) >> 8,
            0x04: test_pattern & 0xFF
        })

        # ADC software reset put its PLL to sleep momentarily. Thus, MMCM needs to be reset as well.
        mmcm_rst()
        
        # Performing Word Align
        perform_bitslip()
        find_edge()
        
        # Printing it once is not enough to check whether the alignment is correct.
        for i in range(100):
            print_adc_channels()

        main_adc_test_mode(spi, False)
        verify_adc_registers(spi, {0x02: 0x11})  # Verify test mode is off
      
        # FIXME: AFE Gain 1x is not functional on that batch of fast servo under development
        enable_adc_afe(ch1_x10=1, ch2_x10=1)

    finally:
        spi.close()

if __name__ == "__main__":
    configure_ltc2195()