#
# Authored by: Robert Metchev / Chips & Scripts (rmetchev@ieee.org)
#
# CERN Open Hardware Licence Version 2 - Permissive
#
# Copyright (C) 2024 Robert Metchev
#

import numpy as np
import sys, os

from dct_aan import aan_scale_factors_2d
#import dct


qt_luma = np.array([
    [16,11,10,16,24,40,51,61],
    [12,12,14,19,26,48,60,55],
    [14,13,16,24,40,57,69,56],
    [14,17,22,29,51,87,80,62],
    [18,22,37,56,68,109,103,77],
    [24,35,55,64,81,104,113,92],
    [49,64,78,87,103,121,120,101],
    [72,92,95,98,112,100,103,99]])

qt_chroma = np.array([
    [17,18,24,47,99,99,99,99],
    [18,21,26,66,99,99,99,99],
    [24,26,56,99,99,99,99,99],
    [47,66,99,99,99,99,99,99],
    [99,99,99,99,99,99,99,99],
    [99,99,99,99,99,99,99,99],
    [99,99,99,99,99,99,99,99],
    [99,99,99,99,99,99,99,99]])
                         

def qt_scale(qt, qf=50):
    """
    Q-table scaler
    Usage: 
        qt_scale(qt_luma, qf=100) / aan_scale_factors_2d
    """    
    if qf < 50:
        scale = 5000/qf
    else:
        scale = 200 - 2*qf  # 2 - qf/50

    t = np.floor((scale*qt + 50) / 100); 
    t[t < 1] = 1   # Prevent divide by 0 error 
    t[t > 255] = 255   # Prevent overflow
    return t.astype(int)

def qt_scale_log(qt, qf_log=0):
    """
    Q-table scaler logarithmic qf=50 -> qfl=0
    Usage: 
        qt_scale(qt_luma, qf_log=0) / aan_scale_factors_2d
    """
    scale = 2**qf_log
    
    # reverse calcuclate qf for reference
    if qf_log >= 0:
        qf = 50/scale
    else:
        qf = 100 - 50*scale
    
    t = np.floor(scale*qt + 0.5); 
    t[t < 1] = 1   # Prevent divide by 0 error 
    t[t > 255] = 255   # Prevent overflow
    return t.astype(int)


# From https://medium.com/100-days-of-algorithms/day-63-zig-zag-51a41127f31
def zig_zag_index(k, n=8):
    # upper side of interval
    if k >= n * (n + 1) // 2:
        i, j = zig_zag_index(n * n - 1 - k, n)
        return n - 1 - i, n - 1 - j
    # lower side of interval
    i = int((np.sqrt(1 + 8 * k) - 1) / 2)
    j = k - i * (i + 1) // 2
    return (j, i - j) if i & 1 else (i - j, j)

# From https://medium.com/100-days-of-algorithms/day-63-zig-zag-51a41127f31
def zig_zag_value(i, j, n=8):
    # upper side of interval
    if i + j >= n:
        return n * n - 1 - zig_zag_value(n - 1 - i, n - 1 - j, n)
    # lower side of interval
    k = (i + j) * (i + j + 1) // 2
    return k + i if (i + j) & 1 else k + j
    
    
def zig_zag_array():
    M = np.empty((8, 8), dtype=int)
    for i in range(8):
        for j in range(8):
            M[i, j] = zig_zag_value(i, j)
    return M
        
def de_zig_zag_array():
    M = np.empty((64), dtype=int)
    for i in range(8):
        for j in range(8):
            M[zig_zag_value(i, j)] = 8*i+j
    return M
        
    
    
    
class QTables:
    def __init__(self, qf=50, qf_log=None):
        self.aan_scale = aan_scale_factors_2d
        self.bits = 13

        self.qf = qf
        self.qf_log = qf_log

        self.qt_luma = qt_scale(qt_luma, self.qf)
        self.qt_luma_aan =  self.qt_luma / aan_scale_factors_2d
        self.qt_luma_aan_factors = 1/self.qt_luma_aan
        self.qt_luma_aan_factors_bin = np.floor((2**(self.bits - 1))/self.qt_luma_aan + 0.5).astype(int)

        self.qt_chroma = qt_scale(qt_chroma, self.qf)
        self.qt_chroma_aan = self.qt_chroma  / aan_scale_factors_2d
        self.qt_chroma_aan_factors = 1/self.qt_chroma_aan
        self.qt_chroma_aan_factors_bin = np.floor((2**(self.bits - 1))/self.qt_chroma_aan + 0.5).astype(int)

    def get_vlog(self):
        x = []
        for y in range(8)[::-1]:
            x.append( ','.join([f'{i:5d}' for i in self.qt_chroma_aan_factors_bin[y][::-1]]))
        for y in range(8)[::-1]:
            x.append( ','.join([f'{i:5d}' for i in self.qt_luma_aan_factors_bin[y][::-1]]))
            
        x = ',\n'.join(x)    

        print (f'// chroma + luma ROMs autogenerated by {os.path.basename(__file__)}')
        print (f'initial mem = {{\n{x}\n}};')

    def quantize_luma(self, m, sel='scipy'):
        if sel == 'aan':
            return np.round(m/self.qt_luma_aan)
        return np.round(m/self.qt_luma)

    def quantize_chroma(self, m, sel='scipy'):
        if sel == 'aan':
            return np.round(m/self.qt_chroma_aan)
        return np.round(m/self.qt_chroma)
    



def get_qt_info(qf):
    qt = QTables(qf)
    print('-' * 79)
    print(f'Luma AAN scaling =\n{qt.aan_scale}\n')
    
    print('-' * 79)
    print(f'QF = {qt.qf}')
    print(f'Bits = {qt.bits}\n')

    print('-' * 79)
    print(f'Luma Q-table =\n{qt_luma}\n')
    print(f'Luma Q-table scaled =\n{qt.qt_luma}\n')
    print(f'Luma Q-table AAN adjusted =\n{qt.qt_luma_aan}\n')
    print(f'Luma Q-table AAN factors =\n{qt.qt_luma_aan_factors}\n')
    print(f'Luma Q-table AAN factors, {qt.bits}-bit =\n{qt.qt_luma_aan_factors_bin}\n')


    print('-' * 79)
    print(f'Chroma Q-table =\n{qt_chroma}\n')
    print(f'Chroma Q-table scaled =\n{qt.qt_chroma}\n')
    print(f'Chroma Q-table AAN adjusted =\n{qt.qt_chroma_aan}\n')
    print(f'Chroma Q-table AAN factors =\n{qt.qt_chroma_aan_factors}\n')
    print(f'Chroma Q-table AAN factors, {qt.bits}-bit =\n{qt.qt_chroma_aan_factors_bin}\n')
    

def get_qf_info():
    print('-' * 79)
    for n in range(-8,8):
        q = 50/2**n
        r = (200 - 100*2**n)/2
        p = q if n > 0 else r
        p = round(p, 1) #if abs(n)==7 else int(p + .5)
        print(f'QF = {p}  scale = 2^{n}')


if __name__ == '__main__':
    qf = 50
    #get_qt_info(qf)
    #get_qf_info()

    qt = QTables(qf)
    
    
    #get_qt_info(qf)
    
    de_zig_zag_index = de_zig_zag_array().flatten()
    
    print('`ifndef __QUANT_TABLES_VH__')
    print('`define __QUANT_TABLES_VH__')
    print(f'// Autogenerated by {os.path.basename(__file__)}')
#    #print('logic[12:0] qt[100:1][1:0][63:0];')
#    #print('logic[25:0] qt[100:1][1:0][32:0];')
#    print('`define QT(h, i, j) qt``h``[i][j]')
#            
#    
#    for qf in range(1,101):
#        qt = QTables(qf)
#        print(f'// QF = {qf}')
#        #print(f'logic[12:0] qt{qf}[1:0][63:0];')
#        print(f'logic[25:0] qt{qf}[1:0][31:0];')
#        for plane in range(2):
#            p = 'Luma' if plane == 0 else 'Chroma'
#            t = qt.qt_luma_aan_factors_bin if plane == 0 else qt.qt_chroma_aan_factors_bin
#            t = t.flatten()[de_zig_zag_index]
#            print(f'// QF = {qf}, {p}')
#            for c in range(32):
#                #print(f'always_comb qt[{qf}][{plane}][{c}] = {t[c]};' )
#                #print(f'always_comb qt{qf}[{plane}][{c}] = {t[c]};' )
#                print(f'always_comb qt{qf}[{plane}][{c}] = 26\'h {((t[2*c + 1] & 0x1FFF) << 13) | (t[2*c] & 0x1FFF):07x};' )


    print('`define QT(h, i) QT``h``_INITVAL_``i``')


    for qf in range(1,101):
        qt = QTables(qf)
        print(f'// QF = {qf}')
        i = 0
        for plane in range(2):
            p = 'Luma' if plane == 0 else 'Chroma'
            t = qt.qt_luma_aan_factors_bin if plane == 0 else qt.qt_chroma_aan_factors_bin
            t = t.flatten()[de_zig_zag_index]
            print(f'// QF = {qf}, {p}')
            #print(t)

            while(len(t)):
                d = t[:16]
                t = t[16:]
        
                m = 0
                for j, vi in enumerate(d):
                    v =  int(vi)
                    #print(j, f'{v:x}')
                    
                    m |= (((v & 0x1ff) | (((v >> 9) & 0x1ff) << 10)) << j*20)
                print(f'localparam QT{qf}_INITVAL_{i:x} = "0x{m:080X}";')
                i += 1




    print('`endif // __QUANT_TABLES_VH__')
