def keeping_solutions(solved):
    """

    :param solved: This is a solution of an equation with multiplicities
    :return: A list of all the solutions

    EXAMPLES:
    sol = solve(-a1^2*x^5 + a2*a3*x^5 + a2*x^4 == 0, x, multiplicities=True, solution_dict=True)
    keeping_solutions(solved)
    [a2/(a1^2 - a2*a3), 0]

    """
    x = var('x')
    solutions = solved[0]
    sols = []
    for solution in solutions:
        sols.append(solution[x])
    return sols


def keeping_multiplicities(solved):
    """

    :param solved: This is a solution of an equation with multiplicities
    :return: A list of all the multiplicities

    EXAMPLES:
    sol = solve(-a1^2*x^5 + a2*a3*x^5 + a2*x^4 == 0, x, multiplicities=True, solution_dict=True)
    keeping_multiplicities(solved)
    [1,4]
    """
    return solved[1]


def minors_solutions_multiplicities(minors):
    """

    :param minors: all the minors of specific dimension of a matrix
    :return: the solutions of all the minors and all the multiplicities for each solution
    """
    solutions_minors = []
    multiplicities_minors = []
    for minor in minors:
        sol_minors = solve(minor == 0, x, multiplicities=True, solution_dict=True)  # solves each minor equation
        solutions_minor = keeping_solutions(sol_minors)  # keeps solutions
        multiplicities_minor = keeping_multiplicities(sol_minors)  # keeps multiplicities
        solutions_minors.append(solutions_minor)
        multiplicities_minors.append(multiplicities_minor)
    return [solutions_minors, multiplicities_minors]


def min_multiplicity_minors(test_sol, minor_sols, minor_mult):
    """

    :param test_sol: the test solution, integer/ symbolic
    :param minor_sols: a list of all the minor solutions, of given dimension
    :param minor_mult: a list of all the multiplicities of the minor solutions, of given dimension
    :return: the minimum multiplicity of test_sol in minor_sols, if it exists, otherwise None
    """
    mult_needed = []
    for j in range(len(minor_sols)):
        if test_sol in minor_sols[j]:  # checks that test_sol is a solution of the minor
            mult_needed.append(minor_mult[j][minor_sols[j].index(test_sol)])  # appends the multiplicity of test_sol
            # in the minor
    if mult_needed:
        return min(mult_needed)
    else:
        return None


def sol_in_minors(test_sol, minor_sol_list, minor_list):
    """

    :param test_sol: the test solution, integer/ symbolic
    :param minor_sol_list: a list of all the minor solutions, of given dimension
    :param minor_list: a list of all the minors, of given dimension
    :return: True or False if the solution belongs in every element of the solution list
    """
    t = 0
    d = 0
    for i in range(len(minor_sol_list)):
        if minor_list[i] != 0:  # we only consider minors that are non-zero
            t = t + 1
            if test_sol in minor_sol_list[i]:  # checks if the solution is a solution of the minor
                d = d + 1
    return t == d


def seg_symbol(mata, matb, n):
    """

    :param mata: an nxn symmetric matrix
    :param matb: an nxn symmetric matrix
    :param n: the dimension of the square matrices
    :return: the Segre symbol of the pencil (mata, matb)

    EXAMPLES:
    F = Matrix([[1,0,0,0,0], [0,1, 0,0,0], [0, 0,1,0,0], [0,0,0,a2,a1], [0,0,0,a1,a3]])
    G = Matrix([[0,0,0,0,0], [0,0, 0,0,0], [0, 0,0,0,0], [0,0,0,0,0], [0,0,0,0,1]])
    seg_symbol(F,G, 5)
    [1, (1, 1, 1, 1)]

    F = Matrix([[0,0, 0,a2], [0,0,a1,a3], [0,a1,a4,a5], [a2,a3,a5,a6] ])
    G = Matrix([[0,0, 0,0], [0,0,0,0], [0,0,b0,b1], [0,0,b1,b2] ])
    seg_symbol(G,F, 4)
    [(2, 2)]
    """
    segre_symbol = []
    x = var('x')
    if mata.det() != 0:
        mat = x * mata + matb
    elif matb.det() != 0:
        mat = mata + x * matb
    else:
        print("Error")
        return None
    all_minors = []
    for k in range(1, n):
        all_minors.append(mat.minors(n-k))  # generates all n-1 up to 1-minors of mat
    all_minors_sols = []
    all_minors_muls = []
    for m in range(len(all_minors)):
        sols_muls = minors_solutions_multiplicities(all_minors[m])  # solves the minor equations for each n-k minor
        all_minors_sols.append(sols_muls[0])
        all_minors_muls.append(sols_muls[1])
    solver = solve(mat.det() == 0, x, multiplicities=True, solution_dict=True)  # solves determinant polynomial equation
    all_solutions = keeping_solutions(solver)
    all_multiplicities = keeping_multiplicities(solver)
    for i in range(len(all_multiplicities)):
        multiplicity = all_multiplicities[i]
        if multiplicity > 1:
            test_solution = all_solutions[i]
            print(test_solution)
            test_val = 0
            for m in range(len(all_minors)):
                if sol_in_minors(test_solution, all_minors_sols[m], all_minors[m]):  # tests if the solution
                    # corresponding to the multiplicity is a solution of all the minors of a specific dimension
                    test_val = test_val + 1
            min_mults = [multiplicity]
            if test_val == len(all_minors):  # if x is a solution for all minors gives segre decomposition (1,...,1)
                seg_dec = [1] * test_val
                segre_symbol.append(tuple(seg_dec))
            elif test_val == 0:  # if x is not a solution of all minors of a specific dimension gives segre
                # decomposition [multiplicity]
                segre_symbol.append(multiplicity)
            else:
                for m in range(test_val):
                    min_mults.append(min_multiplicity_minors(test_solution, all_minors_sols[m], all_minors_muls[m]))
                seg_dec_unrefined = [min_mults[i] - min_mults[i + 1] for i in range(len(min_mults) - 1)]
                seg_dec = [x for x in seg_dec_unrefined if x > 0]  # calculates the segre decomposition
                if sum(seg_dec) < multiplicity:
                    seg_dec.append(min_mults[len(min_mults) - 1])
                segre_symbol.append(tuple(seg_dec))
        else:
            segre_symbol.append(1)  # if not a multiple solution, gives 1 in segre symbol
    return segre_symbol



# ****************************************************************************
#       Copyright (C) 2022 THEODOROS PAPAZACHARIOU <tp19470@essex.ac.uk>
#
# 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 2 of the License, or
# (at your option) any later version.
#                  https://www.gnu.org/licenses/
# ****************************************************************************