"""
Data loading and processing functions for Tilt sensors.

Handles loading raw data from database and initial data structuring.
"""

import numpy as np
import logging
from typing import Dict, Any, Tuple, List
from datetime import datetime
from scipy.signal import medfilt
from ..common.database import DatabaseConnection

logger = logging.getLogger(__name__)


def load_tilt_link_hr_data(
    conn: DatabaseConnection,
    control_unit_id: str,
    chain: str,
    initial_date: str,
    initial_time: str,
    node_list: list
) -> Dict[str, Any]:
    """
    Load Tilt Link HR raw data from database.

    Args:
        conn: Database connection
        control_unit_id: Control unit identifier
        chain: Chain identifier
        initial_date: Starting date
        initial_time: Starting time
        node_list: List of node numbers

    Returns:
        Dictionary with timestamps, angle values, temperatures, and control data
    """
    node_type = 'Tilt Link HR V'

    # Get timestamps from first node
    first_node = node_list[0]

    timestamp_query = """
        SELECT Date, Time
        FROM RawDataView
        WHERE UnitName = %s
          AND ToolNameID = %s
          AND NodeType = %s
          AND NodeNum = %s
          AND (
            (Date = %s AND Time >= %s) OR
            (Date > %s)
          )
        ORDER BY Date, Time
    """

    timestamp_results = conn.execute_query(
        timestamp_query,
        (control_unit_id, chain, node_type, str(first_node),
         initial_date, initial_time, initial_date)
    )

    if not timestamp_results:
        logger.warning("No Tilt Link HR data found")
        return {'timestamps': [], 'values': [], 'errors': []}

    timestamps = []
    for row in timestamp_results:
        dt_str = f"{row['Date']} {row['Time']}"
        timestamps.append(dt_str)

    n_timestamps = len(timestamps)
    logger.info(f"Found {n_timestamps} timestamps for Tilt Link HR data")

    # For TLHR: Val0, Val1 = angles X, Y
    # Val2, Val3, Val4 = control values
    # Val5 = temperature
    n_values_per_node = 6
    all_values = np.zeros((n_timestamps, len(node_list) * n_values_per_node))

    for i, node_num in enumerate(node_list):
        data_query = """
            SELECT Val0, Val1, Val2, Val3, Val4, Val5
            FROM RawDataView
            WHERE UnitName = %s
              AND ToolNameID = %s
              AND NodeType = %s
              AND NodeNum = %s
              AND (
                (Date = %s AND Time >= %s) OR
                (Date > %s)
              )
            ORDER BY Date, Time
        """

        node_results = conn.execute_query(
            data_query,
            (control_unit_id, chain, node_type, str(node_num),
             initial_date, initial_time, initial_date)
        )

        col_offset = i * n_values_per_node
        for j, row in enumerate(node_results):
            if j >= n_timestamps:
                break
            all_values[j, col_offset] = float(row['Val0'] or 0)
            all_values[j, col_offset + 1] = float(row['Val1'] or 0)
            all_values[j, col_offset + 2] = float(row['Val2'] or 0)
            all_values[j, col_offset + 3] = float(row['Val3'] or 0)
            all_values[j, col_offset + 4] = float(row['Val4'] or 0)
            all_values[j, col_offset + 5] = float(row['Val5'] or 0)

        # Forward fill missing data
        if len(node_results) < n_timestamps:
            logger.warning(f"Node {node_num} has only {len(node_results)}/{n_timestamps} records")
            last_valid_idx = len(node_results) - 1
            for j in range(len(node_results), n_timestamps):
                all_values[j, col_offset:col_offset+n_values_per_node] = \
                    all_values[last_valid_idx, col_offset:col_offset+n_values_per_node]

    return {
        'timestamps': timestamps,
        'values': all_values,
        'errors': [],
        'n_nodes': len(node_list)
    }


def define_tilt_link_hr_data(
    raw_data: Dict[str, Any],
    n_sensors: int,
    n_despike: int,
    control_unit_id: str,
    chain: str,
    unit_type: str,
    is_new_zero: bool
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """
    Define and structure Tilt Link HR data from raw database records.

    Converts MATLAB defDatiTLHR.m function.

    Args:
        raw_data: Raw data dict from load_tilt_link_hr_data
        n_sensors: Number of sensors
        n_despike: Number of points for despiking
        control_unit_id: Control unit identifier
        chain: Chain identifier
        unit_type: Unit type identifier
        is_new_zero: Whether this is a new zero point

    Returns:
        Tuple of (timestamps, angles, temperature, control_data, errors)
    """
    if not raw_data or not raw_data.get('values') or len(raw_data['values']) == 0:
        logger.warning("No Tilt Link HR data to define")
        return np.array([]), np.array([]), np.array([]), np.array([]), np.array([])

    logger.info("Defining Tilt Link HR data structure")

    timestamps_str = raw_data['timestamps']
    values = raw_data['values']

    n_timestamps = len(timestamps_str)

    # Convert timestamps to numeric
    timestamps = np.array([
        datetime.strptime(ts, "%Y-%m-%d %H:%M:%S").timestamp()
        for ts in timestamps_str
    ])

    # Extract angles, control data, and temperature
    angles = np.zeros((n_timestamps, n_sensors * 2))
    control_data = np.zeros((n_timestamps, n_sensors * 3))
    temperature = np.zeros((n_timestamps, n_sensors))

    for i in range(n_sensors):
        col_offset = i * 6
        angles[:, i * 2] = values[:, col_offset]      # Val0 = angle X
        angles[:, i * 2 + 1] = values[:, col_offset + 1]  # Val1 = angle Y
        control_data[:, i * 3] = values[:, col_offset + 2]  # Val2
        control_data[:, i * 3 + 1] = values[:, col_offset + 3]  # Val3
        control_data[:, i * 3 + 2] = values[:, col_offset + 4]  # Val4
        temperature[:, i] = values[:, col_offset + 5]  # Val5 = temp

    # Handle NaN values
    n_corrections = 0
    for a in range(1, n_timestamps):
        for b in range(angles.shape[1]):
            if np.isnan(angles[a, b]):
                angles[a, b] = angles[a-1, b]
                n_corrections += 1

    if n_corrections > 0:
        logger.info(f"{n_corrections} NaN values corrected in Tilt Link HR data")

    # Special handling for G301 unit type
    if unit_type == 'G301':
        for i in range(n_sensors):
            for ii in range(1, n_timestamps):
                c_idx = i * 3
                a_idx = i * 2
                # Check for specific error pattern
                if (angles[ii, a_idx] == -8191 and angles[ii, a_idx + 1] == 0 and
                    control_data[ii, c_idx] == 0 and
                    control_data[ii, c_idx + 1] == 0 and
                    control_data[ii, c_idx + 2] == 0):
                    # Copy previous values
                    angles[ii, a_idx:a_idx + 2] = angles[ii-1, a_idx:a_idx + 2]
                    temperature[ii, i] = temperature[ii-1, i]

    # Despiking using median filter
    if n_despike > n_timestamps:
        n_despike = n_timestamps

    for i in range(n_sensors):
        # Apply median filter to remove outliers
        angles[:, i * 2] = medfilt(angles[:, i * 2], kernel_size=n_despike if n_despike % 2 == 1 else n_despike + 1)
        angles[:, i * 2 + 1] = medfilt(angles[:, i * 2 + 1], kernel_size=n_despike if n_despike % 2 == 1 else n_despike + 1)

    # Check for out-of-range values (ampolle fuori scala)
    angles = handle_out_of_range_angles(
        angles, timestamps, control_unit_id, chain, n_sensors, is_new_zero
    )

    # Check for MEMS misreading (ampolla letta come MEMS)
    errors = np.zeros((n_timestamps, n_sensors * 2))
    for b in range(n_sensors):
        c_idx = b * 3
        a_idx = b * 2
        for a in range(n_timestamps):
            # If all control values are non-zero, sensor is being read incorrectly
            if (control_data[a, c_idx] != 0 and
                control_data[a, c_idx + 1] != 0 and
                control_data[a, c_idx + 2] != 0):
                if a > 0:
                    angles[a, a_idx:a_idx + 2] = angles[a-1, a_idx:a_idx + 2]
                errors[a, a_idx:a_idx + 2] = 1

    logger.info(f"Defined Tilt Link HR data: {n_timestamps} timestamps, {n_sensors} sensors")

    return timestamps, angles, temperature, control_data, errors


def handle_out_of_range_angles(
    angles: np.ndarray,
    timestamps: np.ndarray,
    control_unit_id: str,
    chain: str,
    n_sensors: int,
    is_new_zero: bool
) -> np.ndarray:
    """
    Handle out-of-range angle values (scale wrapping at ±32768).

    Args:
        angles: Angle data array
        timestamps: Timestamp array
        control_unit_id: Control unit identifier
        chain: Chain identifier
        n_sensors: Number of sensors
        is_new_zero: Whether this is a new zero point

    Returns:
        Corrected angle array
    """
    # File to store historical angle data
    from pathlib import Path
    import csv

    ampolle_file = Path(f"{control_unit_id}-{chain}-Ampolle.csv")

    # Load previous data if exists
    previous_data = {}
    if is_new_zero and ampolle_file.exists():
        try:
            with open(ampolle_file, 'r') as f:
                reader = csv.reader(f)
                for row in reader:
                    if len(row) > 0:
                        timestamp = float(row[0]) + 730000  # MATLAB datenum offset
                        values = [float(v) for v in row[1:]]
                        previous_data[timestamp] = values
        except Exception as e:
            logger.warning(f"Could not load previous angle data: {e}")

    # Check for scale wrapping
    n_corrections = 0
    for j in range(len(timestamps)):
        for i in range(n_sensors * 2):
            # Get sign of previous value
            if j == 0 and timestamps[j] in previous_data and i < len(previous_data[timestamps[j]]):
                prev_sign = np.sign(previous_data[timestamps[j]][i])
            elif j > 0:
                prev_sign = np.sign(angles[j-1, i])
            else:
                prev_sign = 0

            curr_sign = np.sign(angles[j, i])

            # If signs differ and magnitude is large, scale has wrapped
            if prev_sign != 0 and curr_sign != prev_sign:
                if abs(angles[j, i]) > 15000:
                    if prev_sign == 1:
                        # Positive scale wrap
                        angles[j, i] = 32768 + (32768 + angles[j, i])
                    elif prev_sign == -1:
                        # Negative scale wrap
                        angles[j, i] = -32768 + (-32768 + angles[j, i])
                    n_corrections += 1

    if n_corrections > 0:
        logger.info(f"{n_corrections} out-of-range angle values corrected")

    # Save current data for next run
    try:
        with open(ampolle_file, 'w', newline='') as f:
            writer = csv.writer(f)
            for j in range(len(timestamps)):
                row = [timestamps[j] - 730000] + list(angles[j, :])
                writer.writerow(row)
    except Exception as e:
        logger.warning(f"Could not save angle data: {e}")

    return angles


def load_biaxial_link_data(
    conn: DatabaseConnection,
    control_unit_id: str,
    chain: str,
    initial_date: str,
    initial_time: str,
    node_list: list
) -> Dict[str, Any]:
    """Load Biaxial Link raw data from database."""
    node_type = 'Biaxial Link'

    first_node = node_list[0]

    timestamp_query = """
        SELECT Date, Time
        FROM RawDataView
        WHERE UnitName = %s
          AND ToolNameID = %s
          AND NodeType = %s
          AND NodeNum = %s
          AND (
            (Date = %s AND Time >= %s) OR
            (Date > %s)
          )
        ORDER BY Date, Time
    """

    timestamp_results = conn.execute_query(
        timestamp_query,
        (control_unit_id, chain, node_type, str(first_node),
         initial_date, initial_time, initial_date)
    )

    if not timestamp_results:
        return {'timestamps': [], 'values': [], 'errors': []}

    timestamps = []
    for row in timestamp_results:
        dt_str = f"{row['Date']} {row['Time']}"
        timestamps.append(dt_str)

    n_timestamps = len(timestamps)

    # BL: Val0, Val1 = biaxial data; Val2 = temperature
    n_values_per_node = 3
    all_values = np.zeros((n_timestamps, len(node_list) * n_values_per_node))

    for i, node_num in enumerate(node_list):
        data_query = """
            SELECT Val0, Val1, Val2
            FROM RawDataView
            WHERE UnitName = %s
              AND ToolNameID = %s
              AND NodeType = %s
              AND NodeNum = %s
              AND (
                (Date = %s AND Time >= %s) OR
                (Date > %s)
              )
            ORDER BY Date, Time
        """

        node_results = conn.execute_query(
            data_query,
            (control_unit_id, chain, node_type, str(node_num),
             initial_date, initial_time, initial_date)
        )

        col_offset = i * n_values_per_node
        for j, row in enumerate(node_results):
            if j >= n_timestamps:
                break
            all_values[j, col_offset] = float(row['Val0'] or 0)
            all_values[j, col_offset + 1] = float(row['Val1'] or 0)
            all_values[j, col_offset + 2] = float(row['Val2'] or 0)

    return {
        'timestamps': timestamps,
        'values': all_values,
        'errors': [],
        'n_nodes': len(node_list)
    }


def define_biaxial_link_data(
    raw_data: Dict[str, Any],
    n_sensors: int,
    n_despike: int
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """
    Define and structure Biaxial Link data.

    Args:
        raw_data: Raw data dict
        n_sensors: Number of sensors
        n_despike: Number of points for despiking

    Returns:
        Tuple of (timestamps, data, temperature, errors)
    """
    if not raw_data or not raw_data.get('values'):
        return np.array([]), np.array([]), np.array([]), np.array([])

    timestamps_str = raw_data['timestamps']
    values = raw_data['values']

    n_timestamps = len(timestamps_str)

    timestamps = np.array([
        datetime.strptime(ts, "%Y-%m-%d %H:%M:%S").timestamp()
        for ts in timestamps_str
    ])

    # Extract biaxial data and temperature
    data = np.zeros((n_timestamps, n_sensors * 2))
    temperature = np.zeros((n_timestamps, n_sensors))

    for i in range(n_sensors):
        col_offset = i * 3
        data[:, i * 2] = values[:, col_offset]
        data[:, i * 2 + 1] = values[:, col_offset + 1]
        temperature[:, i] = values[:, col_offset + 2]

    # Despiking
    if n_despike <= n_timestamps:
        for i in range(n_sensors):
            kernel = n_despike if n_despike % 2 == 1 else n_despike + 1
            data[:, i * 2] = medfilt(data[:, i * 2], kernel_size=kernel)
            data[:, i * 2 + 1] = medfilt(data[:, i * 2 + 1], kernel_size=kernel)

    errors = np.zeros((n_timestamps, n_sensors * 2))

    return timestamps, data, temperature, errors
