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

Handles loading raw data from database and initial data structuring.
Converts MATLAB lettura.m and defDati*.m functions.
"""

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

logger = logging.getLogger(__name__)


def load_rsn_link_data(
    conn: DatabaseConnection,
    control_unit_id: str,
    chain: str,
    initial_date: str,
    initial_time: str,
    node_list: list,
    mems_type: int = 2
) -> Dict[str, Any]:
    """
    Load RSN Link raw data from database.
    
    Converts MATLAB lettura.m RSN Link section.
    """
    node_type = 'RSN Link'
    
    # 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 RSN Link data found")
        return {'timestamps': [], 'values': [], 'errors': []}
    
    # Convert timestamps
    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 RSN Link data")
    
    # Load data for each node
    if mems_type == 2:
        value_columns = 'Val0, Val1, Val2, Val6'  # ax, ay, temp, err
        n_values_per_node = 4
    else:
        value_columns = 'Val0, Val1, Val2, Val3'  # ax, ay, az, temp
        n_values_per_node = 4
    
    all_values = np.zeros((n_timestamps, len(node_list) * n_values_per_node))
    errors = []
    
    for i, node_num in enumerate(node_list):
        data_query = f"""
            SELECT {value_columns}
            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)
        )
        
        if not node_results:
            logger.warning(f"No data for RSN node {node_num}")
            errors.append(f"Node {node_num} does NOT work!")
            continue
        
        # Fill data array
        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)
            if mems_type == 2:
                all_values[j, col_offset + 3] = float(row['Val6'] or 0)
            else:
                all_values[j, col_offset + 3] = float(row['Val3'] or 0)
        
        # Handle missing data at end
        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': errors,
        'n_nodes': len(node_list),
        'mems_type': mems_type
    }


def define_rsn_data(
    mems_type: int,
    raw_data: Dict[str, Any],
    error_data: Any,
    n_sensors: int,
    n_despike: int
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """
    Define and structure RSN data from raw database records.
    
    Converts MATLAB defDatiRSN.m function.
    """
    if not raw_data or not raw_data.get('values') or len(raw_data['values']) == 0:
        logger.warning("No RSN data to define")
        return np.array([]), np.array([]), np.array([]), np.array([])
    
    logger.info("Defining RSN 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 acceleration and temperature
    if mems_type == 2:
        # Freescale 2-axis
        n_axes = 2
        acceleration = np.zeros((n_timestamps, n_sensors * n_axes))
        temperature = np.zeros((n_timestamps, n_sensors))
        
        for i in range(n_sensors):
            col_offset = i * 4
            acceleration[:, i * 2] = values[:, col_offset]
            acceleration[:, i * 2 + 1] = values[:, col_offset + 1]
            temperature[:, i] = values[:, col_offset + 2]
    else:
        # 3-axis MEMS
        n_axes = 3
        acceleration = np.zeros((n_timestamps, n_sensors * n_axes))
        temperature = np.zeros((n_timestamps, n_sensors))
        
        for i in range(n_sensors):
            col_offset = i * 4
            acceleration[:, i * 3] = values[:, col_offset]
            acceleration[:, i * 3 + 1] = values[:, col_offset + 1]
            acceleration[:, i * 3 + 2] = values[:, col_offset + 2]
            temperature[:, i] = values[:, col_offset + 3]
    
    # Error flags
    errors = np.zeros((n_timestamps, n_sensors * 4))
    
    logger.info(f"Defined RSN data: {n_timestamps} timestamps, {n_sensors} sensors")
    
    return timestamps, acceleration, temperature, errors
