from problems import *
from algorithms import run_tos
import timeit
import numpy as np

# .. Generate synthetic data ..
np.random.seed(1)

n_iterations = 1000

# Run L1 test
problem = L1Impainting()
gamma_0 = 1.0
def wrapper():
    run_tos(problem, gamma_0, f"runs/TOS_L1.pkl", n_iterations, "nonsmooth", TrueImage=image_orig)
    # NOTE: We do not use the 'TrueImage' in the algorithm. We use it only to evaluate PSNR of our estimates.
print(f"started L1")
execution_time = timeit.timeit(wrapper, number=1)
print(f"L1 done, took {execution_time}")

# Run L2 test
problem = L2Impainting()
gamma_0 = 250.0
def wrapper():
    run_tos(problem, gamma_0, f"runs/TOS_L2.pkl", n_iterations, "nonsmooth", TrueImage=image_orig)
    # NOTE: We do not use the 'TrueImage' in the algorithm. We use it only to evaluate PSNR of our estimates.
print(f"started L2")
execution_time = timeit.timeit(wrapper, number=1)
print(f"L2 done, took {execution_time}")

# Run L2Sq test
problem = L2SqImpainting()
gamma_0 = 1.0
def wrapper():
    run_tos(problem, gamma_0, f"runs/TOS_L2SQ.pkl", n_iterations, "smooth", TrueImage=image_orig)
    # NOTE: We do not use the 'TrueImage' in the algorithm. We use it only to evaluate PSNR of our estimates.
print(f"started L2Sq TOS")
execution_time = timeit.timeit(wrapper, number=1)
print(f"L2Sq TOS done, took {execution_time}")

