import argparse
import gzip
import logging
import pathlib
import pickle
from time import time

import matplotlib.pyplot as plt
import numpy as np
from scipy import interpolate

from .olaflow import OFModel

parser = argparse.ArgumentParser(description="Post-process olaflow results")
parser.add_argument("-v", "--verbose", action="count", default=0)
parser.add_argument(
    "-i", "--input", type=pathlib.Path, help="Olaflow output directory", required=True
)
parser.add_argument(
    "-o",
    "--output",
    type=pathlib.Path,
    help="Output directory for pickled data",
    required=True,
)
parser.add_argument(
    "-z", "--compress", action="store_true", help="Enable gzip compression"
)
args = parser.parse_args()

logging.basicConfig(level=max((10, 20 - 10 * args.verbose)))
log = logging.getLogger("ola_post")

log.info("Starting pickling")
out = args.output
out.mkdir(parents=True, exist_ok=True)

olaflow_root = args.input
model = OFModel(olaflow_root)
model.read_mesh()
model.read_time()

model.read_field_all("alpha.water")
model.read_field_all("porosity")
model.read_field_all("p")
model.read_field_all("p_rgh")
model.read_field_all("U")

model.read_post("graphUniform", "alpha.water")
model.read_post("graphUniform", "U")
model.read_post("graphUniform2", "alpha.water")
model.read_post("graphUniform2", "U")

with (
    gzip.open(out.joinpath("pickle.gz"), "wb")
    if args.compress
    else out.joinpath("pickle").open("wb")
) as f:
    pickle.dump(model, f)