import os
import pandas as pd
import requests
from PIL import Image
from io import BytesIO
from sklearn.preprocessing import LabelEncoder
from sklearn.model_selection import train_test_split
import numpy as np
import tensorflow as tf
from tensorflow.keras import layers, models
from tensorflow.keras.preprocessing.image import img_to_array
import gc

entity_unit_map = {
    'width': {'centimetre', 'foot', 'inch', 'metre', 'millimetre', 'yard'},
    'depth': {'centimetre', 'foot', 'inch', 'metre', 'millimetre', 'yard'},
    'height': {'centimetre', 'foot', 'inch', 'metre', 'millimetre', 'yard'},
    'item_weight': {'gram', 'kilogram', 'microgram', 'milligram', 'ounce', 'pound', 'ton'},
    'maximum_weight_recommendation': {'gram', 'kilogram', 'microgram', 'milligram', 'ounce', 'pound', 'ton'},
    'voltage': {'kilovolt', 'millivolt', 'volt'},
    'wattage': {'kilowatt', 'watt'},
    'item_volume': {'centilitre', 'cubic foot', 'cubic inch', 'cup', 'decilitre', 'fluid ounce', 'gallon', 'imperial gallon', 'litre', 'microlitre', 'millilitre', 'pint', 'quart'}
}

def extract_value_and_unit(entity_value):
    if isinstance(entity_value, str):
        parts = entity_value.split()
        if len(parts) > 0:
            try:
                value = float(parts[0])
            except ValueError:
                value = None
            unit = parts[1] if len(parts) > 1 else ""
            return value, unit
    return None, None

def encode_unit(unit, entity_name):
    if entity_name in entity_unit_map:
        allowed_unit_set = entity_unit_map[entity_name]
        if unit in allowed_unit_set:
            return unit
    return None

def load_image_from_url(image_url, index):
    print(f"Start processing image {index + 1} from URL: {image_url}")
    try:
        response = requests.get(image_url, timeout=5)
        img = Image.open(BytesIO(response.content))
        img = img.resize((224, 224))  
        img_array = img_to_array(img) / 255.0   
        print(f"Successfully processed image {index + 1}")
        return img_array
    except Exception as e:
        print(f"Error loading image from {image_url}: {e}")
        return None

def data_generator(df, batch_size=32):
    while True:
        for start in range(0, len(df), batch_size):
            end = min(start + batch_size, len(df))
            batch_df = df.iloc[start:end]
            
            X_images = []
            valid_indices = []
            for i, url in enumerate(batch_df['image_link']):
                img_array = load_image_from_url(url, i)
                if img_array is not None and img_array.shape == (224, 224, 3):
                    X_images.append(img_array)
                    valid_indices.append(i)
                else:
                    print(f"Skipping invalid image at index {i}")
            
            if len(X_images) == 0:
                continue   
            
            X_images = np.stack(X_images)
            X_entity_names = batch_df['entity_name_encoded'].values[valid_indices]
            X_entity_names = np.expand_dims(X_entity_names, axis=-1)
            y_value = batch_df['numeric_value'].values[valid_indices]
            y_unit = batch_df['unit_encoded'].values[valid_indices]
            
            yield (X_images, X_entity_names), {'entity_value': y_value, 'unit_output': y_unit}
            
            del X_images, X_entity_names, y_value, y_unit
            gc.collect()

def get_dataset(df, batch_size):
    output_signature = (
        (
            tf.TensorSpec(shape=(None, 224, 224, 3), dtype=tf.float32),
            tf.TensorSpec(shape=(None, 1), dtype=tf.int32)
        ),
        {
            'entity_value': tf.TensorSpec(shape=(None,), dtype=tf.float32),
            'unit_output': tf.TensorSpec(shape=(None,), dtype=tf.int32)
        }
    )
    
    return tf.data.Dataset.from_generator(
        lambda: data_generator(df, batch_size),
        output_signature=output_signature
    )

def build_model(num_entity_labels, num_units):
    input_img = layers.Input(shape=(224, 224, 3), name='image_input')
    base_model = tf.keras.applications.ResNet50(weights='imagenet', include_top=False, input_tensor=input_img)
    x = layers.Flatten()(base_model.output)
    x = layers.Dense(512, activation='relu')(x)
    x = layers.Dropout(0.5)(x)

    input_entity_name = layers.Input(shape=(1,), name='entity_name_input')
    entity_name_embedding = layers.Embedding(num_entity_labels, 10)(input_entity_name)
    entity_name_flatten = layers.Flatten()(entity_name_embedding)

    combined = layers.concatenate([x, entity_name_flatten])

    value_output = layers.Dense(1, name='entity_value')(combined)
    unit_output = layers.Dense(num_units, activation='softmax', name='unit_output')(combined)

    model = models.Model(inputs=[input_img, input_entity_name], outputs=[value_output, unit_output])
    return model

class PrintEpochProgress(tf.keras.callbacks.Callback):
    def on_epoch_end(self, epoch, logs=None):
        print(f"Epoch {epoch + 1} finished, Loss: {logs['loss']:.4f}, Val Loss: {logs['val_loss']:.4f}")

class SaveModelCallback(tf.keras.callbacks.Callback):
    def __init__(self, save_every_n_steps):
        super(SaveModelCallback, self).__init__()
        self.save_every_n_steps = save_every_n_steps
        self.steps = 0

    def on_batch_end(self, batch, logs=None):
        self.steps += 1
        if self.steps % self.save_every_n_steps == 0:
            self.model.save(f'saved_models/model_step_{self.steps}.keras')
            print(f"Model saved at step {self.steps}")

def train_on_large_dataset(train_df, val_df, batch_size=16, epochs=10, save_every_n_images=20000):
    num_entity_labels = len(train_df['entity_name_encoded'].unique())
    num_unit_labels = len(train_df['unit_encoded'].unique())
    
    model = build_model(num_entity_labels, num_unit_labels)
    model.compile(optimizer='adam', 
                  loss={'entity_value': 'mse', 'unit_output': 'sparse_categorical_crossentropy'}, 
                  metrics={'entity_value': 'mae', 'unit_output': 'accuracy'})

    train_dataset = get_dataset(train_df, batch_size)
    val_dataset = get_dataset(val_df, batch_size)

    steps_per_epoch = len(train_df) // batch_size
    validation_steps = len(val_df) // batch_size

    save_every_n_steps = save_every_n_images // batch_size

    early_stopping = tf.keras.callbacks.EarlyStopping(patience=3)
    save_callback = SaveModelCallback(save_every_n_steps)
    
    history = model.fit(
        train_dataset,
        steps_per_epoch=steps_per_epoch,
        epochs=epochs,
        validation_data=val_dataset,
        validation_steps=validation_steps,
        callbacks=[early_stopping, save_callback, PrintEpochProgress()]
    )
    
    return model, history

if __name__ == "__main__":
    chunk_size = 10000   
    chunks = pd.read_csv('dataset/train.csv', chunksize=chunk_size)
    
    le_entity_name = LabelEncoder()
    le_units = LabelEncoder()
    
    processed_chunks = []
    for chunk in chunks:
        chunk[['numeric_value', 'unit']] = chunk['entity_value'].apply(lambda x: pd.Series(extract_value_and_unit(x)))
        chunk = chunk.dropna(subset=['numeric_value', 'unit'])
        chunk['filtered_unit'] = chunk.apply(lambda row: encode_unit(row['unit'], row['entity_name']), axis=1)
        chunk = chunk.dropna(subset=['filtered_unit'])
        
        chunk['entity_name_encoded'] = le_entity_name.fit_transform(chunk['entity_name'])
        chunk['unit_encoded'] = le_units.fit_transform(chunk['filtered_unit'])
        
        processed_chunks.append(chunk[['image_link', 'entity_name_encoded', 'numeric_value', 'unit_encoded']])
        
        del chunk
        gc.collect()
    
    train_df = pd.concat(processed_chunks)
    del processed_chunks
    gc.collect()
    
    train_df, val_df = train_test_split(train_df, test_size=0.2, random_state=42)
    
    model, history = train_on_large_dataset(train_df, val_df, batch_size=8, epochs=5, save_every_n_images=20000)

    model.save('final_model.keras')
    print("Training completed. Final model saved as 'final_model.keras'")

    def post_process_predictions(pred_value, pred_unit):
        unit = le_units.inverse_transform


i m running this on my RTX3050 4gb ram, but it is giving error of overflowing ram