import pytest
import pandas as pd
import numpy as np
from omegaconf import OmegaConf
from unittest.mock import MagicMock
import os
import tempfile
import json

# Adjust the import path based on your project structure
from gru_sac_predictor.src.pipeline_stages.sequence_creation import create_sequences_fold
from gru_sac_predictor.src.io_manager import IOManager # Adjust path if needed

# --- Test Fixtures ---

@pytest.fixture
def sample_data_with_imputed():
    """Creates sample X and y dataframes with a 'bar_imputed' column."""
    dates = pd.to_datetime(pd.date_range('2023-01-01', periods=20, freq='T'))
    lookback = 5
    n_features_orig = 3
    n_samples = len(dates)

    # Features (including bar_imputed)
    X_data = pd.DataFrame(
        np.random.randn(n_samples, n_features_orig), 
        index=dates, 
        columns=[f'feat_{i}' for i in range(n_features_orig)]
    )
    # Add bar_imputed column - mark some bars as imputed
    imputed_flags = np.zeros(n_samples, dtype=bool)
    imputed_flags[2] = True # Imputed within first potential sequence
    imputed_flags[8] = True # Imputed within a later potential sequence
    imputed_flags[15] = True # Imputed near the end
    X_data['bar_imputed'] = imputed_flags

    # Targets (mu and dir3)
    y_data = pd.DataFrame({
        'mu': np.random.randn(n_samples),
        'dir3': [list(row) for row in np.eye(3)[np.random.randint(0, 3, n_samples)]] # Example one-hot
    }, index=dates)

    return X_data, y_data

@pytest.fixture
def base_config():
    """Creates a base OmegaConf config for testing sequence creation."""
    conf = OmegaConf.create({
        'gru': {
            'lookback': 5,
            'use_ternary': True, # Matches sample_data_with_imputed
            'drop_imputed_sequences': True # Default to True for testing dropping
        },
        # Add other necessary sections if needed
    })
    return conf

@pytest.fixture
def mock_io_manager():
    """Creates a mock IOManager for testing artefact saving."""
    with tempfile.TemporaryDirectory() as tmpdir:
        mock_io = MagicMock(spec=IOManager)
        mock_io.results_dir = tmpdir
        saved_jsons = {}
        def mock_save_json(data, filename, **kwargs):
            filepath = os.path.join(tmpdir, filename)
            saved_jsons[filename] = data
            with open(filepath, 'w') as f:
                json.dump(data, f, **kwargs)
        mock_io.save_json.side_effect = mock_save_json
        mock_io.get_artifact_path.side_effect = lambda filename: os.path.join(tmpdir, filename)
        mock_io._saved_jsons = saved_jsons 
        yield mock_io

# --- Test Functions ---

def test_sequence_creation_shapes(sample_data_with_imputed, base_config, mock_io_manager):
    X_data, y_data = sample_data_with_imputed
    lookback = base_config.gru.lookback
    n_features = X_data.shape[1]
    n_samples = len(X_data)
    expected_n_seq = n_samples - lookback

    # Test without dropping imputed
    cfg_no_drop = base_config.copy()
    cfg_no_drop.gru.drop_imputed_sequences = False

    X_seq, y_seq_dict, indices, dropped_count = create_sequences_fold(
        X_data=X_data, y_data=y_data, target_names=['mu', 'dir3'], 
        lookback=lookback, name="TestSplit", config=cfg_no_drop, io=mock_io_manager
    )

    assert dropped_count == 0
    assert X_seq is not None
    assert y_seq_dict is not None
    assert indices is not None
    assert X_seq.shape == (expected_n_seq, lookback, n_features)
    assert 'mu' in y_seq_dict and y_seq_dict['mu'].shape == (expected_n_seq,)
    assert 'dir3' in y_seq_dict and y_seq_dict['dir3'].shape == (expected_n_seq, 3)
    assert len(indices) == expected_n_seq
    # Check first target index corresponds to lookback-th original index
    assert indices[0] == X_data.index[lookback]
    # Check last target index corresponds to last original index
    assert indices[-1] == X_data.index[-1]

def test_sequence_dropping_imputed(sample_data_with_imputed, base_config, mock_io_manager):
    X_data, y_data = sample_data_with_imputed
    lookback = base_config.gru.lookback
    n_samples = len(X_data)
    expected_n_seq_orig = n_samples - lookback

    # Config with dropping enabled (default in fixture)
    cfg_drop = base_config

    X_seq, y_seq_dict, indices, dropped_count = create_sequences_fold(
        X_data=X_data.copy(), y_data=y_data.copy(), target_names=['mu', 'dir3'], 
        lookback=lookback, name="TestDrop", config=cfg_drop, io=mock_io_manager
    )

    assert X_seq is not None
    assert y_seq_dict is not None
    assert indices is not None
    
    # Determine which original sequences should have been dropped
    # A sequence starting at index i uses data from [i, i+lookback-1]
    # The target corresponds to index i+lookback
    # We need to check the imputed flag in the range [i, i+lookback-1] for each potential sequence target index i+lookback
    
    # Original target indices range from index `lookback` to `n_samples - 1`
    should_be_dropped_mask = np.zeros(expected_n_seq_orig, dtype=bool)
    imputed_flags_np = X_data['bar_imputed'].values
    for seq_idx in range(expected_n_seq_orig):
        # The features for this sequence are from original indices [seq_idx, seq_idx + lookback - 1]
        feature_indices_range = slice(seq_idx, seq_idx + lookback)
        if np.any(imputed_flags_np[feature_indices_range]):
            should_be_dropped_mask[seq_idx] = True
            
    expected_dropped_count = np.sum(should_be_dropped_mask)
    expected_remaining_count = expected_n_seq_orig - expected_dropped_count

    assert dropped_count == expected_dropped_count
    assert X_seq.shape[0] == expected_remaining_count
    assert y_seq_dict['mu'].shape[0] == expected_remaining_count
    assert y_seq_dict['dir3'].shape[0] == expected_remaining_count
    assert len(indices) == expected_remaining_count
    
    # Check that the remaining indices are correct (weren't marked for dropping)
    original_indices = X_data.index[lookback:]
    expected_remaining_indices = original_indices[~should_be_dropped_mask]
    pd.testing.assert_index_equal(indices, expected_remaining_indices)

    # Check artifact saving
    assert 'imputed_sequence_summary_testdrop.json' in mock_io_manager._saved_jsons
    report_data = mock_io_manager._saved_jsons['imputed_sequence_summary_testdrop.json']
    assert report_data['total_sequences_generated'] == expected_n_seq_orig
    assert report_data['sequences_dropped_imputed'] == expected_dropped_count
    assert report_data['sequences_remaining'] == expected_remaining_count

def test_sequence_creation_no_imputed_col(sample_data_with_imputed, base_config, mock_io_manager):
    X_data, y_data = sample_data_with_imputed
    X_data_no_imputed = X_data.drop(columns=['bar_imputed'])
    lookback = base_config.gru.lookback

    with pytest.raises(SystemExit) as excinfo:
        create_sequences_fold(
            X_data=X_data_no_imputed, y_data=y_data, target_names=['mu', 'dir3'], 
            lookback=lookback, name="TestNoImputedCol", config=base_config, io=mock_io_manager
        )
    assert "'bar_imputed' column missing" in str(excinfo.value)

def test_sequence_creation_insufficient_data(sample_data_with_imputed, base_config, mock_io_manager):
    X_data, y_data = sample_data_with_imputed
    lookback = base_config.gru.lookback
    # Create data shorter than lookback
    X_short = X_data.iloc[:lookback-1]
    y_short = y_data.iloc[:lookback-1]

    X_seq, y_seq_dict, indices, dropped_count = create_sequences_fold(
        X_data=X_short, y_data=y_short, target_names=['mu', 'dir3'], 
        lookback=lookback, name="TestShort", config=base_config, io=mock_io_manager
    )

    assert X_seq is None
    assert y_seq_dict is None
    assert indices is None
    assert dropped_count == 0 