# 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 (
    CH0_HIGH_WORD_ADDR,
    CH0_LOW_WORD_ADDR,
    CH1_HIGH_WORD_ADDR,
    CH1_LOW_WORD_ADDR,
    CTRL_ADDR,
    MAP_MASK,
    PAGESIZE,
    write_to_memory,
    read_from_memory
)

# /dev/spidev2.0  <=> spidev<BUS>.<DEVICE>
MAIN_DAC_BUS = 2
MAIN_DAC_DEVICE = 0
DAC_VERSION = 0x0A


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 soft_reset(spi):
    spi_write(spi, 0x00, 0x10)  # Software reset
    spi_write(spi, 0x00, 0x00)  # Release software reset
    spi_read(spi, 0x00)         # Read reset address (necessary for reset to take effect)

def check_version(spi):
    version = spi_read(spi, 0x1F)
    print(f"DAC version: 0x{version:02X}")
    return version == DAC_VERSION

def configure_dac(spi):
    power_down_reg = spi_read(spi, 0x01)
    spi_write(spi, 0x01, power_down_reg & ~(1 << 0))  # Clear EXTREF bit for internal reference
    spi_write(spi, 0x0D, 0x00)      # Set RREF to 10 kΩ for 1.0V reference
    spi_write(spi, 0x04, 0xA0)      # Enable on-chip IRSET (1.6 kΩ for 20mA output)
    spi_write(spi, 0x07, 0xA0)      # Enable on-chip QRSET (1.6 kΩ for 20mA output)
    spi_write(spi, 0x05, 0x00)      # Disable internal IRCML
    spi_write(spi, 0x08, 0x00)      # Disable internal QRCML
    spi_write(spi, 0x02, 0xB4)      # Enable 2's complement, IFirst: True, IRising: True, DCI_EN: Enabled
    spi_write(spi, 0x14, 0x00)      
    spi_write(spi, 0x14, 0x08)      # Trigger the retimer to reacquire the clock relationship
    spi_write(spi, 0x14, 0x00)

def dac_self_calibration(spi):
    spi_write(spi, 0x12, 0x00)      # Reset calibration status
    spi_write(spi, 0x0E, 0x08)      # Enable calibration clock, default divide ratio
    spi_write(spi, 0x0E, 0x38)      # CALSELI = 1, CALSELQ = 1, CALCLK = 1
    spi_write(spi, 0x12, 0x10)      # Set CALEN bit

    while True:
        status = spi_read(spi, 0x0F)
        if status & 0xC0 == 0xC0:   # Both CALSTATI and CALSTATQ are 1
            break
        time.sleep(0.01)

    spi_write(spi, 0x12, 0x00)      # Clear calibration bits
    spi_write(spi, 0x0E, 0x30)      # Keep CALSELI and CALSELQ set, clear CALCLK
    print("DAC self-calibration completed")

def manual_override(enable=True):
    reg_contents = read_from_memory(CTRL_ADDR, 1)[0]
    print(f"REG contents: 0b{reg_contents:03b}")
    to_write = reg_contents | 0b1 if enable else reg_contents & 0b110
    write_to_memory(CTRL_ADDR, to_write)

def power_down_afe(channel, power_down=True):
    assert channel in (0, 1)

    bitmask = 1 << (channel + 1) & 0b111
    reg_contents = read_from_memory(CTRL_ADDR, 1)[0]
    value = (1 if power_down else 0) << (channel + 1)
    reg_contents &= ~bitmask

    to_write = reg_contents | value
    write_to_memory(CTRL_ADDR, to_write)
    reg_contents = read_from_memory(CTRL_ADDR, 1)[0]
    print(f"REG contents: 0b{reg_contents:03b}")

def set_dac_output(value):
    value = min(value, 0x3FFF)
    low_word = value & 0xFF
    high_word = (value >> 8) & 0x3F

    write_to_memory(CH0_HIGH_WORD_ADDR, high_word)
    write_to_memory(CH0_LOW_WORD_ADDR, low_word)
    write_to_memory(CH1_HIGH_WORD_ADDR, high_word)
    write_to_memory(CH1_LOW_WORD_ADDR, low_word)
    print(f"DAC output set to: 0x{value:04X}")

def check_clk_relationship(spi):
    clkmode_reg = spi_read(spi, 0x14)
    print(f"CLKMODE reg: 0x{clkmode_reg:02X}")
    if clkmode_reg & 0b00010000:
        print("Clock relationship is not found")
        return False
    else:
        print("Clock relationship is found")
        return True

def configure_ad9117():
    spi = spidev.SpiDev()
    spi.open(MAIN_DAC_BUS, MAIN_DAC_DEVICE)
    spi.max_speed_hz = 5000
    spi.mode = 0b00  # CPOL = 0 CPHA = 0
    spi.cshigh = False

    try:
        soft_reset(spi)
        if not check_version(spi):
            print("Unrecognized DAC version")
            return False

        power_down_afe(0, True)
        power_down_afe(1, True)

        configure_dac(spi)
        check_clk_relationship(spi)
        dac_self_calibration(spi)

        # Enable DAC outputs
        spi_write(spi, 0x01, spi_read(spi, 0x01) & ~((1 << 4) | (1 << 3)))

        power_down_afe(0, False)
        power_down_afe(1, False)
        manual_override(False)

        print("AD9117 configuration completed successfully")
        return True

    except Exception as e:
        print(f"Error configuring AD9117: {e}")
        return False

    finally:
        spi.close()

if __name__ == "__main__":
    configure_ad9117()