import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
import pickle

def preprocess_dataset(file_path):
    df = pd.read_csv(file_path, parse_dates=[0])

    result = {}
    for ticker in df['tck_iem_cd'].unique():
        ticker_data = df[df['tck_iem_cd'] == ticker]
        result[ticker] = ticker_data[['gts_iem_ong_pr', 'gts_iem_end_pr', 'gts_iem_hi_pr', 'gts_iem_low_pr']].values.tolist()

    return result


def encode_data(data):
    encoded_data = {}
    for ticker, sequences in data.items():
        encoded_sequences = []
        for sequence in sequences:
            encoded_sequence = []
            for day in sequence:
                open_price, close_price, high_price, low_price = day
                
                # Calculate percentage change (multiply by 100 for percentage representation)
                body_size = 100 * ((close_price - open_price) / open_price)
                upper_shadow = 100 * ((high_price - max(open_price, close_price)) / max(open_price, close_price))
                lower_shadow = 100 * ((min(open_price, close_price) - low_price) / min(open_price, close_price))
                
                encoded_sequence.append([body_size, upper_shadow, lower_shadow])
            encoded_sequences.append(encoded_sequence)
        encoded_data[ticker] = encoded_sequences
    return encoded_data



def create_labels(data):
    labels = {}
    for ticker, sequences in data.items():
        sequence_labels = []
        for i, sequence in enumerate(sequences[:-1]):  # Excluding the last sequence since we don't have the next day data for it
            current_end_price = sequence[-1][1]
            next_end_price = sequences[i+1][-1][1]
            price_change = (next_end_price - current_end_price) / current_end_price
            if price_change >= 0:
                label = 0
            else:
                label = 1
            sequence_labels.append(label)
    
        labels[ticker] = sequence_labels
    return labels


def chunk_sequences(data, window_size=5):
    """
    Chunks the data into sequences of a given window size.
    """
    chunks = []
    for i in range(len(data) - window_size + 1):
        chunks.append(data[i:i+window_size])
    return chunks


def normalize_chunk(chunk):
    """
    Normalizes a given chunk using the first day's start price.
    If the price went up or down by a percentage from the first day's price, 
    that percentage is added to or subtracted from 100.
    """
    first_day_start_price = chunk[0][0]
    normalized_chunk = []

    for day in chunk:
        normalized_day = [
            100 * (day[i] / first_day_start_price) if i == 0 and day_idx == 0 else 
            (100 * (day[i] / first_day_start_price) if day[i] != 0 else 0)
            for day_idx, i in enumerate(range(4))
        ]
        normalized_chunk.append(normalized_day)

    # Setting first day's start price to 100
    normalized_chunk[0][0] = 100
    
    return normalized_chunk


def original_chunk(chunk):
    """
    This function will just return the chunk as it is without performing any normalization.
    """
    return chunk
    

def preprocess_data_for_training(data, window_size=5):
    # Existing functions: chunk_sequences, original_chunk, encode_data, create_labels

    training_data = []

    for ticker, prices in data.items():
        chunks = chunk_sequences(prices, window_size)
        normalized_chunks = [normalize_chunk(chunk) for chunk in chunks]
        original_chunks = [original_chunk(chunk) for chunk in chunks]
        encoded_chunks = encode_data({'temp': original_chunks})['temp']  # Using a temp key to fit the function's format
        
        # Combine original and encoded data
        combined_chunks = [
            [day + encoded_day for day, encoded_day in zip(orig_chunk, encoded_chunk)] 
            for orig_chunk, encoded_chunk in zip(normalized_chunks, encoded_chunks)
        ]

        # Create labels for these chunks
        labels_for_chunks = create_labels({'temp': original_chunks})['temp']

        # Add to the training data
        for combined_chunk, label in zip(combined_chunks, labels_for_chunks):
            training_data.append((combined_chunk, label))

    return training_data
    

file_path = "NASDAQ_DT_FC_STK_QUT.csv"
preprocessed_data = preprocess_dataset(file_path)
training_data = preprocess_data_for_training(preprocessed_data)

# Splitting the data into training and test datasets
train_data, test_data = train_test_split(training_data, test_size=0.2, random_state=42)

# Save train_data
with open('train_data.pkl', 'wb') as f:
    pickle.dump(train_data, f)

# Save test_data
with open('test_data.pkl', 'wb') as f:
    pickle.dump(test_data, f)