import sympy
import scipy.special
import scipy.sparse
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

def skew_symmetric(degree, p_row, p_col):
    if p_col < p_row:
        return p_col, p_row, -1

def zero_diagonal_for_uneven_degree(degree, p_row, p_col):
    if degree % 2 == 1 and p_row == p_col:
        return 0

def gradient(degree, v_row, monom):
    if degree == 1:
        factor = sum(v_row==e for e in monom) + 1

        if monom[-1] < v_row:
            n_v_row = monom[-1]
            n_monom = sorted(monom + (v_row,), reverse=True)

        if v_row <= monom[-1]:
            n_v_row = v_row
            n_monom = monom

        return n_v_row, n_monom, factor

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)

nabla_ha = init_poly_vector(
    degrees=(1,),
    re_index_func_2=gradient,
)
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))

eq = init_equation(
    terms = [(JR, nabla_ha), (JRa, nabla_ha), (JRa, nabla_h), (mg, u)],
    n_var = 2,
    # terms = [(JR, nabla_ha)],
    # n_var = 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(binom(total+2-1, 2))

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

# # mat = init_poly_matrix(
# #     subs={0: {(0, 0): 1, (1, 0): 1, (0, 1): 1, (1, 1): 1}}, 
# # )

# # vec = init_poly_vector(
# #     degrees=(1,),
# #     re_index_func_2=gradient,
# # )

# eq = init_equation(
#     terms = [(mat, vec)],
#     n_var = 2,
# )

eq_tuples, offset_dict, n_param = eq.create()

print(f'{list(offset_dict.values())=}')
print(f'{n_param=}')

# mapping from row index to entry in eq_tuples
rows_to_eq = list(set(key for eq_tuple_degree in eq_tuples.values() for key in eq_tuple_degree.keys()))
eq_to_rows = {eq: idx for idx, eq in enumerate(rows_to_eq)}

print(f'n_eq = {len(rows_to_eq)}')

# # mapping from col index to entry in eq_tuples
# cols_to_var = list(set(key for eq in eq_tuples[degree].values() for key in eq.keys()))
# var_to_cols = {var: idx for idx, var in enumerate(cols_to_var)}

def gen_matrix_per_degree():
    for degree, eq_tuple_degree in eq_tuples.items():

        if 0 < degree:
            def gen_coords(eq_tuple_degree=eq_tuple_degree):
                for eq, var_dict in eq_tuple_degree.items():
                    zero_degree_val = eq_tuples[0][eq][0]
                    row = eq_to_rows[eq]
                    for var, value in var_dict.items():
                        # col = variable_to_index(n_param, var)
                        # col = var_to_cols[var]
                        col = var
                        yield row, col, value / zero_degree_val

            row, col, data = list(zip(*gen_coords()))

            n_col = int(binom(n_param+degree-1, degree))
            n_row = int(max(row) + 1)
            sparse_matrix = scipy.sparse.coo_matrix((data, (row, col)), dtype=np.float64, shape=(n_row, n_col)).toarray()

            yield degree, sparse_matrix

result = dict(gen_matrix_per_degree())

print(result[1].shape)