mapping.py

"""
Модуль для маппинга колонок загруженного датасета на колонки тренировочного датасета
используя Llama через Together API
"""

import pandas as pd
import os
import re
from together import Together


def convert_coordinates_to_degrees(value):
    """
    Конвертирует координаты из формата HMS/DMS в градусы
    Примеры: '07h29m25.85s' -> 112.357708 degrees
             '45d30m15.5s' -> 45.504306 degrees
    """
    if pd.isna(value) or isinstance(value, (int, float)):
        return value
    
    value_str = str(value).strip()
    
    # Формат HMS (часы:минуты:секунды) для RA
    hms_match = re.match(r'(\d+)h(\d+)m([\d.]+)s?', value_str)
    if hms_match:
        hours = float(hms_match.group(1))
        minutes = float(hms_match.group(2))
        seconds = float(hms_match.group(3))
        return hours * 15 + minutes * 0.25 + seconds * 0.00416667  # 1h = 15°, 1m = 0.25°, 1s = 0.00416667°
    
    # Формат DMS (градусы:минуты:секунды) для DEC
    dms_match = re.match(r'([+-]?)(\d+)d(\d+)m([\d.]+)s?', value_str)
    if dms_match:
        sign = -1 if dms_match.group(1) == '-' else 1
        degrees = float(dms_match.group(2))
        minutes = float(dms_match.group(3))
        seconds = float(dms_match.group(4))
        return sign * (degrees + minutes / 60 + seconds / 3600)
    
    # Если не распознали формат, возвращаем NaN
    return float('nan')


class ColumnMapper:
    def __init__(self, api_key: str):
        """
        Initialize column mapper
        
        Args:
            api_key: API key for Together AI
        """
        self.client = Together(api_key=api_key)
        
        # Built-in synonym dictionary (fallback) - significantly expanded
        self.known_synonyms = {
            # Orbital period
            'pl_orbper': 'koi_period',
            'orbital_period': 'koi_period',
            'period': 'koi_period',
            'pl_orbpererr1': 'koi_period_err1',
            'pl_orbpererr2': 'koi_period_err2',
            'pl_orbpererr': 'koi_period_err1',
            
            # Transit time/epoch
            'pl_tranmid': 'koi_time0bk',
            'transit_time': 'koi_time0bk',
            'time0': 'koi_time0bk',
            'epoch': 'koi_time0bk',
            'pl_tranmiderr1': 'koi_time0bk_err1',
            'pl_tranmiderr2': 'koi_time0bk_err2',
            
            # Transit duration
            'pl_trandur': 'koi_duration',
            'pl_trandurh': 'koi_duration',
            'transit_duration': 'koi_duration',
            'duration': 'koi_duration',
            'pl_trandurerr1': 'koi_duration_err1',
            'pl_trandurerr2': 'koi_duration_err2',
            
            # Transit depth
            'pl_trandep': 'koi_depth',
            'transit_depth': 'koi_depth',
            'depth': 'koi_depth',
            'pl_trandeperr1': 'koi_depth_err1',
            'pl_trandeperr2': 'koi_depth_err2',
            
            # Planet radius
            'pl_rade': 'koi_prad',
            'pl_radj': 'koi_prad',
            'planet_radius': 'koi_prad',
            'radius': 'koi_prad',
            'pl_radeerr1': 'koi_prad_err1',
            'pl_radeerr2': 'koi_prad_err2',
            'pl_radjerr1': 'koi_prad_err1',
            'pl_radjerr2': 'koi_prad_err2',
            
            # Insolation flux
            'pl_insol': 'koi_insol',
            'insolation': 'koi_insol',
            'insol': 'koi_insol',
            'pl_insolerr1': 'koi_insol_err1',
            'pl_insolerr2': 'koi_insol_err2',
            
            # Equilibrium temperature
            'pl_eqt': 'koi_teq',
            'equilibrium_temp': 'koi_teq',
            'teq': 'koi_teq',
            'pl_eqterr1': 'koi_teq_err1',
            'pl_eqterr2': 'koi_teq_err2',
            
            # Stellar effective temperature
            'st_teff': 'koi_steff',
            'stellar_teff': 'koi_steff',
            'star_temp': 'koi_steff',
            'teff': 'koi_steff',
            'st_tefferr1': 'koi_steff_err1',
            'st_tefferr2': 'koi_steff_err2',
            
            # Stellar surface gravity
            'st_logg': 'koi_slogg',
            'stellar_logg': 'koi_slogg',
            'surface_gravity': 'koi_slogg',
            'logg': 'koi_slogg',
            'st_loggerr1': 'koi_slogg_err1',
            'st_loggerr2': 'koi_slogg_err2',
            
            # Stellar radius
            'st_rad': 'koi_srad',
            'stellar_radius': 'koi_srad',
            'star_radius': 'koi_srad',
            'st_raderr1': 'koi_srad_err1',
            'st_raderr2': 'koi_srad_err2',
            
            # Stellar mass
            'st_mass': 'koi_smass',
            'stellar_mass': 'koi_smass',
            'st_masserr1': 'koi_smass_err1',
            'st_masserr2': 'koi_smass_err2',
            
            # Kepler magnitude
            'sy_kepmag': 'koi_kepmag',
            'kepmag': 'koi_kepmag',
            'kep_mag': 'koi_kepmag',
            'sy_kepmaglim': 'koi_kepmag',
            
            # Coordinates
            'ra': 'ra',
            'ra_deg': 'ra',
            'rastr': 'ra',
            'dec': 'dec',
            'dec_deg': 'dec',
            'decstr': 'dec',
            
            # Model SNR
            'koi_model_snr': 'koi_model_snr',
            'snr': 'koi_model_snr',
            
            # Impact parameter
            'pl_imppar': 'koi_impact',
            'impact': 'koi_impact',
            'impact_parameter': 'koi_impact',
            
            # Additional mappings for error columns
            'koi_period_err': 'koi_period_err1',
            'koi_time0bk_err': 'koi_time0bk_err1',
            'koi_duration_err': 'koi_duration_err1',
            'koi_depth_err': 'koi_depth_err1',
            'koi_prad_err': 'koi_prad_err1',
            'koi_teq_err': 'koi_teq_err1',
            'koi_insol_err': 'koi_insol_err1',
            'koi_steff_err': 'koi_steff_err1',
            'koi_slogg_err': 'koi_slogg_err1',
            'koi_srad_err': 'koi_srad_err1',
            'koi_smass_err': 'koi_smass_err1',
        }
        
    def get_column_mapping(self, source_columns: list, target_columns: list) -> dict:
        """
        Получает маппинг между колонками источника и целевыми колонками
        используя LLM
        
        Args:
            source_columns: Список колонок загруженного датасета
            target_columns: Список колонок тренировочного датасета
            
        Returns:
            Словарь маппинга {source_column: target_column}
        """
        # Словарь известных синонимов для точного маппинга
        known_mappings = """
Common column name mappings (NASA Exoplanet Archive):
- pl_orbper, orbital_period, period → koi_period (Orbital Period in days)
- pl_tranmid, transit_time, time0 → koi_time0bk (Transit Epoch in BJD)
- pl_trandur, pl_trandurh, transit_duration → koi_duration (Transit Duration in hours)
- pl_trandep, transit_depth, depth → koi_depth (Transit Depth in ppm)
- pl_rade, planet_radius, radius → koi_prad (Planetary Radius in Earth radii)
- pl_insol, insolation, insol → koi_insol (Insolation Flux in Earth flux)
- pl_eqt, equilibrium_temp, teq → koi_teq (Equilibrium Temperature in K)
- st_teff, stellar_teff, star_temp → koi_steff (Stellar Effective Temperature in K)
- st_logg, stellar_logg, surface_gravity → koi_slogg (Stellar Surface Gravity in log10(cm/s^2))
- st_rad, stellar_radius, star_radius → koi_srad (Stellar Radius in Solar radii)
- st_mass, stellar_mass, star_mass → koi_smass (Stellar Mass in Solar masses)
- ra, ra_deg → ra (Right Ascension in degrees)
- dec, dec_deg → dec (Declination in degrees)
- pl_bmassj, planet_mass → koi_prad (use radius if mass not available)
- sy_dist, distance → koi_steff (stellar distance - related to stellar properties)
"""
        
        prompt = f"""You are an expert in NASA Exoplanet Archive data mapping. Map column names from a source dataset to Kepler/KOI target dataset columns.

{known_mappings}

Source columns:
{source_columns}

Target columns:
{target_columns}

CRITICAL INSTRUCTIONS:
1. Use the known mappings above as your PRIMARY reference
2. Match columns based on physical meaning (e.g., "pl_orbper" = orbital period = "koi_period")
3. Common prefixes: "pl_" = planet property, "st_" = stellar property, "koi_" = KOI property
4. If exact match exists in known mappings, USE IT
5. Only map columns with clear semantic similarity
6. Return ONLY a Python dictionary: {{"source": "target", ...}}
7. NO markdown, NO explanations, NO code blocks - just the dictionary

Example: {{"pl_orbper": "koi_period", "st_teff": "koi_steff", "ra": "ra"}}

Mapping:"""

        response = self.client.chat.completions.create(
            model="meta-llama/Llama-3.3-70B-Instruct-Turbo",
            messages=[{"role": "user", "content": prompt}],
            temperature=0.1,
            max_tokens=2000
        )
        
        mapping_str = response.choices[0].message.content.strip()
        
        # Очистка ответа от возможных markdown блоков
        if "```" in mapping_str:
            mapping_str = mapping_str.split("```")[1]
            if mapping_str.startswith("python"):
                mapping_str = mapping_str[6:]
            mapping_str = mapping_str.strip()
        
        # Преобразование строки в словарь
        try:
            mapping = eval(mapping_str)
            if not isinstance(mapping, dict):
                raise ValueError("Response is not a dictionary")
        except Exception as e:
            print(f"Error parsing mapping: {e}")
            print(f"Raw response: {mapping_str}")
            # Возвращаем пустой маппинг в случае ошибки
            mapping = {}
        
        # Supplement mapping with known synonyms (fallback)
        # Check source columns that were not mapped by Llama
        unmapped_sources = [col for col in source_columns if col not in mapping]
        
        for src_col in unmapped_sources:
            src_lower = src_col.lower()
            
            # Check exact match with known synonyms
            if src_lower in self.known_synonyms:
                target = self.known_synonyms[src_lower]
                if target in target_columns:
                    mapping[src_col] = target
                    continue
            
            # Check for partial matches (more sophisticated)
            # Remove common prefixes/suffixes for comparison
            src_clean = src_lower.replace('pl_', '').replace('st_', '').replace('sy_', '').replace('koi_', '')
            
            for known_src, known_tgt in self.known_synonyms.items():
                known_clean = known_src.replace('pl_', '').replace('st_', '').replace('sy_', '').replace('koi_', '')
                
                # Check if core part matches
                if src_clean == known_clean or known_clean in src_clean or src_clean in known_clean:
                    if known_tgt in target_columns:
                        mapping[src_col] = known_tgt
                        break
            
            # If still not mapped, try fuzzy matching on target columns
            if src_col not in mapping:
                for tgt_col in target_columns:
                    tgt_clean = tgt_col.replace('koi_', '')
                    # Check if source contains target name
                    if tgt_clean in src_lower or src_clean == tgt_clean:
                        mapping[src_col] = tgt_col
                        break
            
        return mapping
    
    def apply_mapping(self, df: pd.DataFrame, mapping: dict) -> pd.DataFrame:
        """
        Применяет маппинг к датафрейму
        
        Args:
            df: Исходный датафрейм
            mapping: Словарь маппинга
            
        Returns:
            Датафрейм с переименованными колонками
        """
        # Переименовываем только те колонки, которые есть в маппинге
        df_mapped = df.copy()
        
        # Проверяем какие колонки из маппинга действительно есть в датафрейме
        valid_mapping = {k: v for k, v in mapping.items() if k in df.columns}
        
        if valid_mapping:
            df_mapped = df_mapped.rename(columns=valid_mapping)
        
        # Ensure all columns are properly flattened and converted to numeric where possible
        for col in df_mapped.columns:
            try:
                # Get the column as a Series
                col_data = df_mapped[col]
                
                # Check if it's actually a Series (not a DataFrame)
                if not isinstance(col_data, pd.Series):
                    continue
                
                # Check if column has object dtype or might contain complex data
                if col_data.dtype == 'object':
                    try:
                        # Try to convert to numeric
                        df_mapped[col] = pd.to_numeric(col_data, errors='coerce')
                    except:
                        pass
                
                # Ensure column is 1D
                if hasattr(col_data, 'values'):
                    col_values = col_data.values
                    if len(col_values.shape) > 1:
                        # Flatten multi-dimensional arrays
                        df_mapped[col] = col_values.flatten()[:len(df_mapped)]
            except Exception as e:
                # Skip problematic columns
                continue
            
        return df_mapped
    
    def map_dataset(self, uploaded_df: pd.DataFrame, target_columns: list) -> tuple:
        """
        Полный процесс маппинга датасета
        
        Args:
            uploaded_df: Загруженный датафрейм
            target_columns: Список колонок тренировочного датасета
            
        Returns:
            Кортеж (mapped_dataframe, mapping_dict, info_message)
        """
        # Копируем датафрейм чтобы не изменять оригинал
        df_work = uploaded_df.copy()
        
        # Clean up column names - remove extra spaces, special characters
        df_work.columns = df_work.columns.str.strip()
        
        # Handle any multi-dimensional columns before mapping
        for col in df_work.columns:
            if df_work[col].dtype == 'object':
                # Check if column contains complex structures
                first_val = df_work[col].dropna().iloc[0] if len(df_work[col].dropna()) > 0 else None
                
                if isinstance(first_val, (list, tuple)):
                    # Flatten lists/tuples - take first element
                    df_work[col] = df_work[col].apply(
                        lambda x: x[0] if isinstance(x, (list, tuple)) and len(x) > 0 else (x if not isinstance(x, (list, tuple)) else None)
                    )
                elif isinstance(first_val, str):
                    # Try to convert string representations of numbers
                    try:
                        df_work[col] = pd.to_numeric(df_work[col], errors='ignore')
                    except:
                        pass
        
        # Конвертируем координаты в градусы если они в текстовом формате
        coord_columns = [col for col in df_work.columns if any(
            keyword in col.lower() for keyword in ['ra', 'dec', 'coord', 'right_ascension', 'declination', 'rastr', 'decstr']
        )]
        
        for col in coord_columns:
            # Check first non-empty value
            first_val = df_work[col].dropna().iloc[0] if len(df_work[col].dropna()) > 0 else None
            if first_val and isinstance(first_val, str) and ('h' in first_val or 'd' in first_val):
                # Convert entire column
                df_work[col] = df_work[col].apply(convert_coordinates_to_degrees)
        
        source_columns = df_work.columns.tolist()
        
        # Get mapping via LLM
        mapping = self.get_column_mapping(source_columns, target_columns)
        
        # Apply mapping
        mapped_df = self.apply_mapping(df_work, mapping)
        
        # Create info message
        if mapping:
            info_msg = f"Successfully mapped {len(mapping)} columns:\n"
            for src, tgt in mapping.items():
                info_msg += f"  * {src} -> {tgt}\n"
        else:
            info_msg = "Warning: No mapping performed - no matches found between columns\n"
            info_msg += f"Source columns: {', '.join(source_columns[:5])}...\n"
        
        # Check which target columns are missing
        missing_cols = set(target_columns) - set(mapped_df.columns)
        if missing_cols:
            info_msg += f"\nWarning: Missing {len(missing_cols)} target columns (will be filled with zeros)\n"
        
        return mapped_df, mapping, info_msg


def load_training_columns(csv_path: str) -> list:
    """
    Load column names from training dataset
    
    Args:
        csv_path: Path to training dataset CSV file
        
    Returns:
        List of column names
    """
    df = pd.read_csv(csv_path, comment='#', nrows=1)
    return df.columns.tolist()