#!/usr/bin/python3
# -*- coding: utf-8 -*-

from challenges.opti.babystep_giantstep import BSGS
from maths.math_lib import CRT
from sage.all import factor


def pohlig_hellman_prime(G, A, E, p, e):
    power = [1] * (e + 1)  # power[k] = pi ** k
    for k in range(1, e + 1):
        power[k] = power[k - 1] * p

    inv = -G
    y = power[e - 1] * G

    x, h, d = [0] * (e + 1), [None] * e, [0] * e

    for k in range(e):
        # compute the k-th digit in base pi of x mod pi^ni
        h[k] = power[e - 1 - k] * (x[k] * inv + A)
        d[k] = BSGS(y, h[k], E(0), p)
        x[k + 1] = x[k] + d[k] * power[k]

    return x[e]


def pohlig_hellman(G, A, E):
    """Pohlig-Hellman algorithm for computing discrete logarithms
    in an elliptic curve `E` whose order is a smooth integer.

        INPUT:
        - `G` must be an element of `E`
        - `A` must be an element of the sub-group generated by `G`

        OUTPUT:
        - `n` such that `n*G = A`
    """

    n = G.order()
    decomposition = factor(n)
    residues = []
    moduli = []

    for pi, ei in decomposition:
        # compute the residue of x mod pi^ni
        modulus = pi ** ei
        ni = n // modulus
        Gi, Ai = ni * G, ni * A
        x = pohlig_hellman_prime(Gi, Ai, E, pi, ei)

        residues.append(x)
        moduli.append(modulus)

    x = CRT(residues, moduli)
    return x