import sympy
import scipy.special
import numpy as np
from polymatrix.polystruct import init_equation, init_poly_matrix, init_poly_vector
from polymatrix.sympyutils import poly_to_data_coord, poly_to_matrix
from polymatrix.utils import variable_to_index

binom = scipy.special.binom

x = sympy.symbols('v_dc, i_d, i_q')
vdc, id, iq = x

g_dc = 3
wl = 106

jr_poly_list = [
    g_dc, -1, g_dc*wl,
    1, 0, wl,
    -g_dc*wl, -wl, 0
]
jr_terms = poly_to_data_coord(list(sympy.poly(p, x) for p in jr_poly_list))
# print(jr_terms)

h = vdc**2/2 + id**2/2 + iq**2/2
dh_poly_list = [
    sympy.diff(h, vdc), sympy.diff(h, id), sympy.diff(h, iq)
]
dh_terms = poly_to_data_coord(list(sympy.poly(p, x) for p in dh_poly_list))
# print(dh_terms)

mg_poly_list = [
    g_dc-id, -iq,
    vdc+1, 0,
    0, vdc+1
]
mg_terms = poly_to_data_coord(list(sympy.poly(p, x) for p in mg_poly_list))
# print(mg_terms)

nabla_h = init_poly_vector(subs=dh_terms)
JR = init_poly_matrix(subs=jr_terms)
mg = init_poly_matrix(subs=mg_terms)

def skew_symmetric(idx1, idx2):
    if idx2 < idx1:
        return idx2, idx1, -1

def zero_diagonal_for_uneven_degree(degree, idx1, idx2):
    if degree % 2 == 1 and idx1 == idx2:
        return 0

    

nabla_ha = init_poly_vector(degrees=(1,2))
JRa = init_poly_matrix(
    degrees=(0,1,2), 
    re_index_func=skew_symmetric, 
    subs_func=zero_diagonal_for_uneven_degree,
)
u = init_poly_vector(degrees=(1,2))

# n_var = 2
# n_param_nabla_ha = sum(n_var*binom(n_var+p-1, p) for p in nabla_ha.degrees)
# n_param_JRa = sum(n_var**2*binom(n_var+p-1, p) for p in JRa.non_zero_degrees)
# n_param_u = sum(n_var*binom(n_var+p-1, p) for p in u.non_zero_degrees)
# total = n_param_nabla_ha+n_param_JRa+n_param_u

# print(f'{n_param_nabla_ha=}')
# print(f'{n_param_JRa=}')
# print(f'{n_param_u=}')
# print(f'{total=}')

# print(n_var*binom(n_var+1-1, 1))
# print(n_var*binom(n_var+2-1, 2))

# print(binom(total+2-1, 2))

# eq = init_equation(
#     # terms = [(JR, nabla_ha), (JRa, nabla_ha), (JRa, nabla_h), (mg, u)],
#     # n_var = 3,
#     terms = [(JR, nabla_ha)],
#     n_var = 2,
# )

# print(eq.create())

mat = init_poly_matrix(degrees=(0,), re_index_func=skew_symmetric)
# vec = init_poly_vector(degrees=(0,1))
vec = init_poly_vector(subs={0: {(0, 0): 1, (1, 0): 1}})

eq = init_equation(
    terms = [(mat, vec)],
    n_var = 2,
)
eq_tuples, offset_dict = eq.create()

# print(eq_tuples)
# print(eq_tuples[1][(1, tuple())])
# print(offset_dict[(mat, 0)])
# print(offset_dict[(vec, 0)])
print(offset_dict)
# # print(variable_to_index(10, (offset_dict[(mat, 0)], offset_dict[(vec, 1)])))