#!/usr/bin/env python3

# TP2 info601
# NOM : Alonzo Church
# groupe L3-info-00101101

import sys
import os
from heapq import heapify, heappush, heappop

from naive_ppm import Image


def find_seam(img):
    oo = float("inf")       # infinite value used as default value for D[s] when s hasn't been seen yet

    D = {(i, 0): img.E[0][i] for i in range(img.width)}
    Pred = {}

    ATraiter = [(img.E[0][i], i, 0) for i in range(img.width)]
    heapify(ATraiter)

    while ATraiter:
        e, x, y = heappop(ATraiter)

        if y == img.height - 1:
            return Pred, x

        ...

        for v in ...:
            ...

    assert False    # We shouldn't ever finish the loop without reaching the last line


########################################################################
########################################################################
########################################################################


def carve(filename, nb_seams=10, mask=None, show_step=None, save_initial_energy=False):
    base, ext = os.path.splitext(filename)

    img = Image(filename, mask)

    if save_initial_energy:
        img.save_E(base + "-energy.ppm")

    savefile = base + "-carved" + ext
    if show_step is not None:
        seamfile_template = base + "-{:03}-seam" + ext
        # savefile_template = base + "-{:03}-carved" + ext
        # maskfile_template = base + "-{:03}-mask" + ext
        # energyfile_template = base + "-{:03}-energy" + ext

    for n in range(nb_seams):
        print(f"searching for seam {n} / {nb_seams}: ", end="", flush=True)
        Pred, x = find_seam(img)
        S = []
        xy = (x, img.height-1)
        e = 0
        while xy is not None:
            S.append(xy[0])
            e += img.E[xy[1]][xy[0]]
            xy = Pred.get(xy)
        assert len(S) == img.height
        S.reverse()

        print(f"found seam with total energy: {e}")

        if show_step is not None and n % show_step == 0:
            img.draw_seam(S)
            img.save(seamfile_template.format(n))
            # img.save_mask(maskfile_template.format(n))
            # img.save_E(energyfile_template.format(n))

        img.carve_seam(S)

        # if show_step is not None and n % show_step == 0:
        #     img.save(savefile_template.format(n))

    print(f"\nsave result into {savefile}\n")
    img.save(savefile)


def int_or_exit(s):
    try:
        return int(s)
    except ValueError:
        print(f"cannot convert '{s}' to an integer")
        sys.exit(3)


if __name__ == "__main__":

    if len(sys.argv) == 1 or len(sys.argv) > 5:
        print(f"usage: {sys.argv[0]} FILE_PPM [MASK_PPM] [NB_SEAMS [STEP]] ")
        sys.exit(1)

    filename = sys.argv[1]
    save_energy = False

    if len(sys.argv) == 2:
        nb_seams = 1
        step = 1
        mask = None
    elif len(sys.argv) == 3:
        if sys.argv[2].endswith(".ppm"):
            mask = sys.argv[2]
            nb_seams = 1
            step = 1
        else:
            mask = None
            nb_seams = int_or_exit(sys.argv[2])
            step = None
    elif len(sys.argv) == 4:
        if sys.argv[2].endswith(".ppm"):
            mask = sys.argv[2]
            nb_seams = int_or_exit(sys.argv[3])
            step = None
        else:
            mask = None
            nb_seams = int_or_exit(sys.argv[2])
            step = int_or_exit(sys.argv[3])
    elif len(sys.argv) == 5:
        if sys.argv[2].endswith(".ppm"):
            mask = sys.argv[2]
            nb_seams = int_or_exit(sys.argv[3])
            step = int_or_exit(sys.argv[4])
        else:
            print(f"usage: {sys.argv[0]} FILE_PPM [MASK_PPM] [NB_SEAMS [STEP]] ")
            sys.exit(1)

    carve(filename, nb_seams, mask=mask, show_step=step, save_initial_energy=save_energy)