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

# System modules
from time import sleep

# Third-party modules
# import gnuplotlib as gp # No longer needed
import numpy as np
from astropy.io import fits
from astropy.time import Time

# Other templates
from ..config.constants import pixscale, temp_fits  # Corrected relative import
from ..devices import cam, foc  # Corrected relative import
from ..utils.analysis import fit_star  # Corrected relative import
from ..utils.logger import log  # Corrected relative import
from ..utils.structure import foc_path  # Corrected relative import
from .basetemplate import BaseTemplate  # Corrected relative import
from .box import Template as Box  # Corrected relative import


def simple_ascii_plot(
        x_values,
        y_values,
        x_label="X",
        y_label="Y",
        title="Plot",
        width=50,
        height=15):
    """
    Generates a simple ASCII plot for the terminal.
    Scales y_values to fit within the height.
    """
    if not len(x_values) or not len(
            y_values) or len(x_values) != len(y_values):
        log.warning("ASCII Plot: Invalid or empty data for plotting.")
        return ""

    plot_str = []
    plot_str.append(f"{title:^{width}}")
    plot_str.append("")  # Empty line

    min_y, max_y = min(y_values), max(y_values)
    y_range = max_y - min_y
    if y_range == 0:  # Avoid division by zero if all y values are the same
        y_range = 1.0

    # Create the plot grid (transposed for easier filling)
    grid = [[' ' for _ in range(width)] for _ in range(height)]

    # Y-axis labels and scale
    y_axis_label_width = 8  # max width for y-axis numbers
    plot_width_effective = width - y_axis_label_width - 1  # -1 for the axis itself

    # Determine x positions (scaled to fit plot_width_effective)
    # This simple version assumes x_values are somewhat evenly spaced or just for labeling
    # For true x-axis scaling, it's more complex for ASCII.
    # We will place points at somewhat regular intervals across the width.

    num_points = len(x_values)
    if num_points == 0:
        return ""

    # Map y_values to plot height
    scaled_y = [int(((y - min_y) / y_range) * (height - 1)) for y in y_values]

    # Create a simple bar-like representation for each point
    # This is a very basic representation.
    # We will map each x_value to a column if possible, or group them.

    # For this simple plot, we'll iterate through y_values and represent them as bars
    # at distinct x positions if possible.

    # Let's make a bar chart where each x_value has its own "bar"
    # The width of each bar will be plot_width_effective / num_points

    bar_segment_width = max(1, plot_width_effective // num_points)

    for i in range(height):
        row_str = ""
        # Y-axis label
        y_val_at_row = max_y - \
            (i * y_range / (height - 1 if height > 1 else 1))
        row_str += f"{y_val_at_row:<{y_axis_label_width}.2f}|"

        for j in range(num_points):
            # Check if the current row corresponds to the bar height for this point
            # scaled_y is 0 at min_y, height-1 at max_y.
            # Grid is top-down, so we need to invert.
            if (height - 1 - scaled_y[j]) <= i:
                row_str += "*" * bar_segment_width
            else:
                row_str += " " * bar_segment_width
        plot_str.append(row_str.rstrip())

    # X-axis (very simplified)
    plot_str.append(
        "-" *
        y_axis_label_width +
        "+" +
        "-" *
        (plot_width_effective))

    # X-axis labels (approximate positions)
    # This is tricky to align perfectly in ASCII.
    # We'll just show min and max x for simplicity or label a few points.
    if num_points > 0:
        x_labels_line = " " * (y_axis_label_width + 1)
        if num_points == 1:
            x_labels_line += f"{x_values[0]:^{plot_width_effective}.1f}"
        elif num_points > 1:
            first_x_label = f"{x_values[0]:.1f}"
            last_x_label = f"{x_values[-1]:.1f}"
            middle_space = plot_width_effective - \
                len(first_x_label) - len(last_x_label)
            if middle_space < 0:
                middle_space = 0
            x_labels_line += first_x_label + " " * middle_space + last_x_label
        plot_str.append(x_labels_line)
        plot_str.append(f"{x_label:^{width}}")

    plot_str.append("")
    plot_str.append(f"{y_label} vs {x_label}")
    return "\n".join(plot_str)


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 400 // binning
            height = params.get("height") or 400 // binning
            repeat = params.get("repeat") or 10
            step = params.get("step") or 10

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

        # for Box
        fixed = {
            "objname": objname,
            "repeat": 1,
            "center": center,
            "width": width,
            "height": height,
        }
        params.update(fixed)

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

        # Getting initial focus
        original_foc = foc.position
        if original_foc is None:
            log.error("Could not get initial focus position. Aborting focus.")
            self.error.append("Initial focus position is None.")
            return

        log.info(f"Initial focus position: {original_foc} µm")

        # starting from half before
        # Corrected: (repeat-1) steps for 'repeat' points
        total_length_focus_scan = step * (repeat - 1)
        start_focus = original_foc - total_length_focus_scan / 2

        log.info(
            f"Starting focus scan from {
                start_focus:.0f} µm to {
                start_focus +
                total_length_focus_scan:.0f} µm in {repeat} steps of {step} µm.")

        # Move to the start of the focus scan
        foc.position = int(round(start_focus))
        # Add a small delay for the focuser to settle if it's a physical device
        sleep(1.0)  # Adjust as needed, or check foc.is_moving if available

        # Preparing empty arrays for values
        m2_arr = np.array([], dtype=float)  # Use float for focus positions
        fwhm_arr = np.array([], dtype=float)
        self.output = []

        # Preparing the focus file path
        now_time = Time.now().isot
        # Ensure foc_path is from ..utils.structure
        file_path = foc_path(now_time)

        # Initial comment about variables
        comment1 = f"Focus procedure of {now_time} for object '{objname}' with filter '{filt}'"
        comment2 = f"Step Focus[µm] FWHM[arcsec] Peak[ADU] BG[ADU] X[px] Y[px]"

        # Write the initial comment if the file doesn't exist
        with open(file_path, "a+") as file:
            log.info(f"Using focus data file: {file_path}")
            file.seek(0)
            if not file.read(1):
                file.write(f"# {comment1}\n")
                file.write(f"# {comment2}\n")

        # Instanciating a Box template
        box_template = Box()

        for rep_idx in range(repeat):
            current_focus_target = int(round(start_focus + rep_idx * step))
            log.info(
                f"Focus step {
                    rep_idx + 1}/{repeat} :: Target M2 Focus: {current_focus_target} µm")

            # Set focus position for this step
            foc.position = current_focus_target
            # Check if focuser is moving (if available) or sleep
            # This is important to ensure the image is taken at the correct
            # focus
            if hasattr(foc, 'is_moving'):
                while foc.is_moving:
                    log.debug("Waiting for focuser to settle...")
                    sleep(0.2)  # Short sleep while checking
                    if self.check_pause_or_abort():
                        return
            else:
                sleep(1.0)  # Generic delay if is_moving is not available

            # Get actual focus position after movement
            actual_m2_pos = foc.position
            if actual_m2_pos is None:
                log.warning(
                    f"Could not read M2 position at step {
                        rep_idx + 1}. Skipping this point.")
                continue

            m2_arr = np.append(m2_arr, actual_m2_pos)

            ################################
            ##### Taking a boxed image #####
            ################################
            # Update exptime and filter in params for the box template, if they
            # can change
            params["exptime"] = exptime
            params["filter"] = filt
            params["binning"] = binning

            box_template.run(params)  # saves temp.fits boxed image
            if self.check_pause_or_abort():
                return

            try:
                data = fits.getdata(temp_fits)
            except FileNotFoundError:
                log.error(
                    f"Focus: {temp_fits} not found after box exposure. Skipping point.")
                # Add NaN for missing FWHM
                fwhm_arr = np.append(fwhm_arr, np.nan)
                continue

            ###############################
            ##### Fitting a 2d Moffat #####
            ###############################
            current_fwhm_arcsec = np.nan  # Default to NaN
            peak_adu, bg_adu, fit_x, fit_y = np.nan, np.nan, np.nan, np.nan

            try:
                # Estimate background from corners if possible, or use a fixed guess
                # For simplicity, using a fixed guess for now
                background_estimation = np.median(
                    data)  # A slightly better guess
                fitted = fit_star(
                    data,
                    model="moffat",
                    background_estimation=background_estimation)
                if fitted and hasattr(
                        fitted,
                        'fwhm') and hasattr(
                        fitted,
                        'xc'):
                    current_fwhm_arcsec = fitted.fwhm * pixscale * binning  # FWHM in arcsec
                    peak_adu = fitted.peak
                    bg_adu = fitted.background
                    fit_x = fitted.xc
                    fit_y = fitted.yc
                    log.info(
                        f"  -> Fit: FWHM={
                            current_fwhm_arcsec:.2f}\", Peak={
                            peak_adu:.0f}, BG={
                            bg_adu:.0f}, X={
                            fit_x:.1f}, Y={
                            fit_y:.1f}")
                else:
                    log.warning(
                        "  -> Fit failed or did not return expected attributes.")
            except Exception as e:
                log.warning(
                    f"  -> Star fitting error at focus {actual_m2_pos} µm: {e}")

            fwhm_arr = np.append(fwhm_arr, current_fwhm_arcsec)

            self.output = {  # Update self.output for potential real-time display/API
                "focus_run": {
                    "m2_positions": m2_arr.tolist(),
                    # Replace NaN for JSON
                    "fwhm_arcsec": np.nan_to_num(fwhm_arr, nan=-1.0).tolist(),
                    "current_point": {
                        "m2": actual_m2_pos,
                        "fwhm": current_fwhm_arcsec if not np.isnan(current_fwhm_arcsec) else -1.0,
                        "peak": peak_adu if not np.isnan(peak_adu) else -1.0,
                        "background": bg_adu if not np.isnan(bg_adu) else -1.0,
                        "fit_x_px": fit_x if not np.isnan(fit_x) else -1.0,
                        "fit_y_px": fit_y if not np.isnan(fit_y) else -1.0,
                    }
                }
            }

            with open(file_path, "a") as file:
                line_data = f"{
                    rep_idx +
                    1:<4d} {
                    actual_m2_pos:<10.1f} {
                    current_fwhm_arcsec:<10.2f} {
                    peak_adu:<9.0f} {
                    bg_adu:<7.0f} {
                        fit_x:<7.1f} {
                            fit_y:<7.1f}"
                file.write(line_data + "\n")

            if self.check_pause_or_abort():
                return

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

        log.info("Focus scan complete. Results:")
        for i in range(len(m2_arr)):
            log.info(f"  Focus: {m2_arr[i]:.1f} µm, FWHM: {fwhm_arr[i]:.2f}\"")

        # Generate and print ASCII plot
        if len(m2_arr) > 1 and len(fwhm_arr) > 1:
            # Filter out NaN values for plotting and finding minimum
            valid_indices = ~np.isnan(fwhm_arr)
            plot_m2 = m2_arr[valid_indices]
            plot_fwhm = fwhm_arr[valid_indices]

            if len(plot_m2) > 0:
                ascii_chart = simple_ascii_plot(
                    plot_m2, plot_fwhm,
                    x_label="Focus M2 (µm)", y_label="FWHM (arcsec)",
                    title=f"Focus Curve - {objname} ({filt})",
                    width=80, height=20
                )
                log.info("\n" + ascii_chart)

                # Find best focus
                if len(plot_fwhm) > 0:
                    min_fwhm_idx = np.argmin(plot_fwhm)
                    best_focus_val = plot_m2[min_fwhm_idx]
                    best_fwhm_val = plot_fwhm[min_fwhm_idx]
                    log.info(
                        f"Optimal focus found at {
                            best_focus_val:.1f} µm with FWHM {
                            best_fwhm_val:.2f}\".")
                    log.info(
                        f"Moving M2 to optimal focus: {
                            best_focus_val:.0f} µm.")
                    foc.position = int(round(best_focus_val))
                else:
                    log.warning(
                        "No valid FWHM data points to determine optimal focus. Restoring original.")
                    foc.position = original_foc
            else:
                log.warning(
                    "No valid data points to plot or determine optimal focus. Restoring original.")
                foc.position = original_foc
        else:
            log.warning(
                "Not enough data points for a plot or to determine optimal focus. Restoring original focus.")
            foc.position = original_foc

        log.info(
            f"Focus procedure finished. M2 position set to: {
                foc.position} µm.")
