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

# System modules
from astropy.io import fits
import gnuplotlib as gp
import numpy as np
from time import sleep

# Custom modules
from devices import cam, foc
from utils.logger import log
from utils.analysis import fit_box2d
from config.constants import temp_fits, pixscale

# Other templates
from templates.basetemplate import BaseTemplate
from templates.box import Template as Box


class Template(BaseTemplate):
    def __init__(self):
        super().__init__()
        self.name = "focus"
        self.description = "Fit boxed images to get optimal focus"

    def content(self, params):

        ########################
        ##### Params check #####
        ########################

        try:
            objname = params.get("object") or "test"
            filt = params["filter"]
            binning = params["binning"]
            exptime = params["exptime"]
            center = params.get("center") or cam.center
            width = params.get("width") or 80//binning  # px
            height = params.get("height") or 80//binning  # ~24''
            repeat = params.get("repeat") or 5
            step = params.get("step") or 20

        except KeyError as e:
            log.error(f"Parameter {e} not found")
            return

        fixed = {"repeat": 1,
                 "frametype": "Light"
                 }

        params.update(fixed)

        #######################################
        ##### Preparing to the focus loop #####
        #######################################

        # Getting initial focus
        original_foc = foc.position

        # Preparing empty arrays for values
        m2 = np.array([], dtype=int)
        fwhm = np.array([])

        # Instanciating a Box template
        box = Box()

        for rep in range(0, repeat):
            log.info(f"Step {rep} of {repeat-1}")

            ################################
            ##### Taking a boxed image #####
            ################################

            box.run(params)  # saves temp.fits boxed image

            img = fits.getdata(temp_fits)

            ###############################
            ##### Fitting a 2d Moffat #####
            ###############################

            fitted = fit_box2d(img, model="moffat")

            ##############################
            ##### Filling the arrays #####
            ##############################

            m2 = np.append(m2, foc.position)
            fwhm = np.append(fwhm, fitted.fwhm*pixscale*binning)

            log.info(f"Found: focus={m2}µm, FWHM={fwhm}'' ")

            ######################
            ##### Moving M2  #####
            ######################

            foc.position += step
            log.warning(f"Moving focus of {step}µm to {foc.position}µm")
            sleep(0.3)

        ##############################
        ##### Showing the result #####
        ##############################

        log.info(f"focus={m2}, fwhm={fwhm}")

        gp.plot((m2, fwhm),
                unset="grid",
                _with="linespoints",
                terminal='dumb 90 30',
                xlabel="x=focus [µm], y=fwhm ['']")

        #################################
        ##### Back to initial focus #####
        #################################

        log.info("Setting back original focus")
        foc.position = original_foc

        # test return
        return {"m2": m2, "fwhm": fwhm}
