import math
import pynq
import pytest
import importlib
import os
import numpy as np
from .mock_devices import MockMemoryMappedDevice


ZYNQ_SLCR_OFFSET = 0xF8000000
ZU_LPD_OFFSET = 0xFF5E0000
ZU_FPD_OFFSET = 0xFD1A0000


class FakeUname:
    def __init__(self, machine):
        self.machine = machine


def be_zynq():
    return FakeUname(pynq.ps.ZYNQ_ARCH)


def be_zu():
    return FakeUname(pynq.ps.ZU_ARCH)


def be_other():
    return FakeUname('Invalid Arch')


# Entries of the form PLL_OFFSET, CLK_REG
ZYNQ_READ_PLLS = {
    'arm_pll': [0x100, 0x0030_0520],
    'ddr_pll': [0x104, 0x0030_0530],
    'io0_pll': [0x108, 0x0030_0500],
    'io1_pll': [0x108, 0x0030_0510],
}

ZYNQ_READ_CLKS = {
    'fclk0_mhz': 0x170,
    'fclk1_mhz': 0x180,
    'fclk2_mhz': 0x190,
    'fclk3_mhz': 0x1A0,
}

ZYNQ_CPU_PLLS = {
    'arm0_pll': [0x100, 0x0200],
    'arm1_pll': [0x100, 0x0210],
    'ddr_pll': [0x104, 0x0220],
    'io_pll': [0x108, 0x0230],
}

ZU_READ_PLLS = {
    'iopll': (1, 0x20, 0x0050300),
    'rpll': (1, 0x30, 0x0050302),
    'dpll': (0, 0x2C, 0x0050303),
}

ZU_READ_CLKS = {
    'fclk0_mhz': 0xc0,
    'fclk1_mhz': 0xc4,
    'fclk2_mhz': 0xc8,
    'fclk3_mhz': 0xcC,
}

ZU_CPU_PLLS = {
    'apll': (0x20, 0x200),
    'dpll': (0x2C, 0x202),
    'vpll': (0x38, 0x203),
}


@pytest.fixture
def setup_zynq(monkeypatch):
    old_arch = pynq.ps.CPU_ARCH
    monkeypatch.setattr(os, 'uname', be_zynq)
    device = MockMemoryMappedDevice('zynq_clocks')
    pynq.Device.active_device = device
    slcr_buffer = device.mmap(ZYNQ_SLCR_OFFSET, 0x200)
    slcr_array = np.frombuffer(slcr_buffer, dtype='u4')
    new_ps = importlib.reload(pynq.ps)
    Clocks = new_ps.Clocks

    yield Clocks, slcr_array

    pynq.Device.active_device = None
    pynq.ps.CPU_ARCH = old_arch


@pytest.fixture
def setup_zu(monkeypatch):
    old_arch = pynq.ps.CPU_ARCH
    monkeypatch.setattr(os, 'uname', be_zu)
    device = MockMemoryMappedDevice('zu_clocks')
    pynq.Device.active_device = device
    lpd_buffer = device.mmap(ZU_LPD_OFFSET, 0x100)
    lpd_array = np.frombuffer(lpd_buffer, dtype='u4')

    fpd_buffer = device.mmap(ZU_FPD_OFFSET, 0x100)
    fpd_array = np.frombuffer(fpd_buffer, dtype='u4')

    new_ps = importlib.reload(pynq.ps)
    Clocks = new_ps.Clocks

    yield Clocks, lpd_array, fpd_array

    pynq.Device.active_device = None
    pynq.ps.CPU_ARCH = old_arch


def split_clk_reg(reg_val):
    source = (reg_val >> 4) & 0x3
    div0 = (reg_val >> 8) & 0x3F
    div1 = (reg_val >> 20) & 0x3F
    return source, div0, div1


def split_zu_reg(reg_val):
    source = (reg_val) & 0x07
    div0 = (reg_val >> 8) & 0x3F
    div1 = (reg_val >> 16) & 0x3F
    return source, div0, div1


@pytest.mark.parametrize('pll_name', ZYNQ_READ_PLLS.keys())
@pytest.mark.parametrize('clk_name', ZYNQ_READ_CLKS.keys())
def test_zynq_pl(setup_zynq, pll_name, clk_name):
    pll_reg, clk_val = ZYNQ_READ_PLLS[pll_name]
    clk_reg = ZYNQ_READ_CLKS[clk_name]
    Clocks, slcr_array = setup_zynq

    slcr_array[pll_reg >> 2] = 30 << 12  # FBDIV of 30 (1500 MHz)
    slcr_array[clk_reg >> 2] = clk_val
    assert getattr(Clocks, clk_name) == 100

    setattr(Clocks, clk_name, 50)
    source, div0, div1 = split_clk_reg(slcr_array[clk_reg >> 2])
    assert div0 * div1 == 30