# 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/>.

# Additional Reference:
# https://github.com/torvalds/linux/blob/master/drivers/clk/clk-si5341.c

import time
from smbus2 import SMBus

BUS_NO = 0
IC_ADDR = 0x74

DEVICE_READY = 0x00FE
PLL_M_DEN = 0x023B
STATUS = 0x000C
STATUS_STICKY = 0x0011

STATUS_LOSREF = 0x04
STATUS_LOL = 0x08


def write_register(bus, address, value):
    page = address >> 8
    register = address & 0xFF
    bus.write_byte_data(IC_ADDR, 0x01, page)
    try:
        bus.write_byte_data(IC_ADDR, register, value)
    except Exception as e:
        raise Exception(f"Write failed 0x{value:02X} at 0x{address:04X}: {e}")

def write_preamble(bus):
    preamble = [
        (0x0B24, 0xC0),
        (0x0B25, 0x00),
        (0x0502, 0x01),
        (0x0505, 0x03),
        (0x0957, 0x17),
        (0x0B4E, 0x1A),
    ]
    for address, value in preamble:
        write_register(bus, address, value)

def write_postamble(bus):
    postamble = [
        (0x001C, 0x01),  # Soft reset
        (0x0B24, 0xC3),
        (0x0B25, 0x02),
    ]
    for address, value in postamble:
        write_register(bus, address, value)

def wait_device_ready(bus):
    for _ in range(15):
        if bus.read_byte_data(IC_ADDR, DEVICE_READY) == 0x0F:
            return True
        time.sleep(0.02)
    return False

def wait_for_lock(bus):
    for _ in range(10):
        status = bus.read_byte_data(IC_ADDR, STATUS)
        if not (status & (STATUS_LOSREF | STATUS_LOL)):
            return True
        time.sleep(0.01)
    return False

def check_pll_status(bus):
    pll_status = bus.read_byte_data(IC_ADDR, 0x0C)
    pll_locked = not (pll_status & STATUS_LOL)
    print(f"PLL {'locked' if pll_locked else 'unlocked'}")
    return pll_locked

def check_los_status(bus):
    los_status = bus.read_byte_data(IC_ADDR, 0x0D)
    xaxb_los = (los_status & 0x10) != 0
    print(f"XA/XB LOS {'asserted' if xaxb_los else 'deasserted'}")
    return not xaxb_los

def configure_si5340():
    with SMBus(BUS_NO) as bus:
        if not wait_device_ready(bus):
            print("Device not ready. Aborting.")
            return

        # Programming sequence from ClockBuilder Pro, default settings
        # to initialize system using XTAL input
        main_config = [
            (0x0006, 0x00),  # TOOL_VERSION
            (0x0007, 0x00),  # Not in datasheet
            (0x0008, 0x00),  # Not in datasheet
            (0x000B, 0x74),  # I2C_ADDR
            (0x0017, 0xD0),  # INT mask (disable interrupts)
            (0x0018, 0xFF),  # INT mask
            (0x0021, 0x0F),  # Select XTAL as input
            (0x0022, 0x00),  # Not in datasheet
            (0x002B, 0x02),  # SPI config
            (0x002C, 0x20),  # LOS enable for XTAL
            (0x002D, 0x00),  # LOS timing
            (0x002E, 0x00),  # LOS trigger (thresholds)
            (0x002F, 0x00),
            (0x0030, 0x00),
            (0x0031, 0x00),
            (0x0032, 0x00),
            (0x0033, 0x00),
            (0x0034, 0x00),
            (0x0035, 0x00),  # LOS trigger (thresholds) end
            (0x0036, 0x00),  # LOS clear (thresholds)
            (0x0037, 0x00),
            (0x0038, 0x00),
            (0x0039, 0x00),
            (0x003A, 0x00),
            (0x003B, 0x00),
            (0x003C, 0x00),
            (0x003D, 0x00),  # LOS clear (thresholds) end
            (0x0041, 0x00),  # LOS0_DIV_SEL
            (0x0042, 0x00),  # LOS1_DIV_SEL
            (0x0043, 0x00),  # LOS2_DIV_SEL
            (0x0044, 0x00),  # LOS3_DIV_SEL
            (0x009E, 0x00),  # LOL_SET_THR
            (0x0102, 0x01),  # Enable outputs
            (0x013F, 0x00),  # OUTX_ALWAYS_ON
            (0x0140, 0x00),  # OUTX_ALWAYS_ON
            (0x0141, 0x40),  # OUT_DIS_LOL_MSK, OUT_DIS_MSK_LOS_PFD
            (0x0202, 0x00),  # XAXB_FREQ_OFFSET (=0)

            # PLL Configuration
            (0x0235, 0x00),  # M_NUM
            (0x0236, 0x00),
            (0x0237, 0x00),
            (0x0238, 0x80),
            (0x0239, 0x89),
            (0x023A, 0x00),
            (0x023B, 0x00),  # M_DEN
            (0x023C, 0x00),
            (0x023D, 0x00),
            (0x023E, 0x80),

            # Synthesizer configuration
            (0x0302, 0x00),  # N0_NUM
            (0x0303, 0x00),
            (0x0304, 0x00),
            (0x0305, 0x00),
            (0x0306, 0x21),
            (0x0307, 0x00),
            (0x0308, 0x00),  # N0_DEN
            (0x0309, 0x00),
            (0x030A, 0x00),
            (0x030B, 0x80),
            (0x030C, 0x01), # N0_UPDATE

            # N1 Configuration (1:1 ratio)
            (0x030D, 0x00),  # N1_NUM
            (0x030E, 0x00),
            (0x030F, 0x00),
            (0x0310, 0x00),
            (0x0311, 0x00),
            (0x0312, 0x01),
            (0x0313, 0x00),  # N1_DEN
            (0x0314, 0x00),
            (0x0315, 0x00),
            (0x0316, 0x01),
            (0x0317, 0x01),  # N1_UPDATE

            # N2 Configuration (1:1 ratio)
            (0x0318, 0x00),  # N2_NUM
            (0x0319, 0x00),
            (0x031A, 0x00),
            (0x031B, 0x00),
            (0x031C, 0x00),
            (0x031D, 0x01),
            (0x031E, 0x00),  # N2_DEN
            (0x031F, 0x00),
            (0x0320, 0x00),
            (0x0321, 0x01),
            (0x0322, 0x01),  # N2_UPDATE

            # N3 Configuration (1:1 ratio)
            (0x0323, 0x00),  # N3_NUM
            (0x0324, 0x00),
            (0x0325, 0x00),
            (0x0326, 0x00),
            (0x0327, 0x00),
            (0x0328, 0x01),
            (0x0329, 0x00),  # N3_DEN
            (0x032A, 0x00),
            (0x032B, 0x00),
            (0x032C, 0x01),
            (0x032D, 0x01),  # N3_UPDATE

            # Output configuration
            (0x0112, 0x06),  # OUT0 config
            (0x0113, 0x09),  # OUT0 format
            (0x0114, 0x3B),  # OUT0 CM/AMPL
            (0x0115, 0x28),  # OUT0 MUX_SEL

            (0x0117, 0x06),  # OUT1 config
            (0x0118, 0x09),  # OUT1 format
            (0x0119, 0x3B),  # OUT1 CM/AMPL
            (0x011A, 0x28),  # OUT1 MUX_SEL

            (0x0126, 0x06),  # OUT2 config
            (0x0127, 0x09),  # OUT2 format
            (0x0128, 0x3B),  # OUT2 CM/AMPL
            (0x0129, 0x28),  # OUT2 MUX_SEL

            (0x012B, 0x06),  # OUT3 config
            (0x012C, 0xCC),  # OUT3 format
            (0x012D, 0x00),  # OUT3 CM/AMPL
            (0x012E, 0x58),  # OUT3 MUX_SEL

            # Miscellaneous configuration
            (0x090E, 0x02),  # XAXB_EXTCLK_EN=0 XAXB_PDNB=1 (use XTAL)
            (0x091C, 0x04),  # ZDM_EN=4 (Normal mode)
            (0x0943, 0x00),  # IO_VDD_SEL
            (0x0949, 0x00),  # IN_EN (disable input clocks)
            (0x094A, 0x00),  # INx_TO_PFD_EN (disabled)
            (0x094E, 0x49),  # REFCLK_HYS_SEL (set by CBPro)
            (0x094F, 0x02),  # Not in datasheet
            (0x095E, 0x00),  # M_INTEGER (set by CBPro)
            (0x0A02, 0x00),  # N_ADD_0P5 (set by CBPro)
            (0x0A03, 0x01),  # N_CLK_TO_OUTX_EN
            (0x0A04, 0x01),  # N_PIBYP
            (0x0A05, 0x01),  # N_PDNB
            (0x0A14, 0x00),  # N0_HIGH_FREQ (set by CBPro)
            (0x0A1A, 0x00),  # N1_HIGH_FREQ (set by CBPro)
            (0x0A20, 0x00),  # N2_HIGH_FREQ (set by CBPro)
            (0x0A26, 0x00),  # N3_HIGH_FREQ (set by CBPro)
            (0x0B44, 0x0F),  # PDIV_ENB (set by CBPro)
            (0x0B4A, 0x0E),  # N_CLK_DIS
            (0x0B57, 0x0E),  # VCO_RESET_CALCODE (set by CBPro)
            (0x0B58, 0x01),  # VCO_RESET_CALCODE (set by CBPro)
        ]

        write_preamble(bus)

        time.sleep(0.3)

        print("Writing main configuration...")
        for address, value in main_config:
            write_register(bus, address, value)
        print("Main configuration written")

        write_postamble(bus)

        if not wait_for_lock(bus):
            print("Error waiting for input clock or PLL lock")
        else:
            print("Input clock present and PLL locked")

        bus.write_byte_data(IC_ADDR, STATUS_STICKY, 0)

        # Final status check
        pll_locked = check_pll_status(bus)
        xaxb_signal_present = check_los_status(bus)

        if not pll_locked:
            print("Error: PLL is not locked")
        elif not xaxb_signal_present:
            print("Error: XA/XB signal is lost")
        else:
            print("Si5340 configuration completed successfully")

if __name__ == "__main__":
    configure_si5340()