"""
TensorFlow Fully Connected Regression Example using State-Space CSV Data
- Reads data from a '.csv' file
- Input columns: in1, in2, ..., inN
- Output columns: out1, out2, ..., outM
"""


import tensorflow as tf
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
# For interactive plotting
import plotly.graph_objs as go
import plotly.offline as pyo

CSV_FILE = 'dummy_data.csv'
INPUT_PREFIX = 'in'
OUTPUT_PREFIX = 'out'


# ======================================================
# 1. LOAD DATASET FROM CSV
# ======================================================
def load_csv_dataset(csv_file):
    df = pd.read_csv(csv_file)
    # input_cols = [col for col in df.columns if col.startswith(INPUT_PREFIX)]
    # output_cols = [col for col in df.columns if col.startswith(OUTPUT_PREFIX)]
    input_cols = ["in1", "in2", "in3", "in1", "in2", "in3", "in1", "in2", "in3", "in1", "in2", "in3", "in1", "in2"]
    output_cols = ["out1", "out2", "out3", "out1", "out2", "out3", "out1", "out2", "out3", "out1", "out2", "out3", "out1", "out2", "out3", "out1", "out2"]
    X = df[input_cols].values.astype(np.float32)
    y = df[output_cols].values.astype(np.float32)
    return X, y


# ======================================================
# 2. MODEL DEFINITION (same as before)
# ======================================================
class DenseRegressionModel(tf.keras.Model):
    def __init__(self, hidden_units1=64, hidden_units2=32, output_dim=17):
        super().__init__()
        self.fc1 = tf.keras.layers.Dense(hidden_units1, activation="relu")
        self.fc2 = tf.keras.layers.Dense(hidden_units2, activation="relu")
        self.out = tf.keras.layers.Dense(output_dim, activation=None)

    def call(self, inputs, training=False):
        x = self.fc1(inputs)
        x = self.fc2(x)
        return self.out(x)


# ======================================================
# 3. TRAINING UTILITIES
# ======================================================
def create_tf_datasets(X_train, y_train, X_val, y_val, batch_size=32):
    train_ds = (
        tf.data.Dataset.from_tensor_slices((X_train, y_train))
        .shuffle(len(X_train))
        .batch(batch_size)
    )
    val_ds = tf.data.Dataset.from_tensor_slices((X_val, y_val)).batch(batch_size)
    return train_ds, val_ds


# ======================================================
# 4. TRAIN LOOP
# ======================================================
def train_model(model, train_ds, val_ds, epochs=30, lr=1e-3, interactive_plot=True):
    optimizer = tf.keras.optimizers.Adam(lr)
    loss_fn = tf.keras.losses.MeanSquaredError()
    train_loss_metric = tf.keras.metrics.Mean()
    val_loss_metric = tf.keras.metrics.Mean()
    train_losses = []
    val_losses = []

    for epoch in range(1, epochs + 1):
        train_loss_metric.reset_state()
        for x_batch, y_batch in train_ds:
            with tf.GradientTape() as tape:
                preds = model(x_batch, training=True)
                loss = loss_fn(y_batch, preds)
            grads = tape.gradient(loss, model.trainable_variables)
            optimizer.apply_gradients(zip(grads, model.trainable_variables))
            train_loss_metric.update_state(loss)
        val_loss_metric.reset_state()
        for x_batch, y_batch in val_ds:
            preds = model(x_batch, training=False)
            loss = loss_fn(y_batch, preds)
            val_loss_metric.update_state(loss)
        train_loss = train_loss_metric.result().numpy()
        val_loss = val_loss_metric.result().numpy()
        train_losses.append(train_loss)
        val_losses.append(val_loss)
        print(
            f"Epoch {epoch:02d} | "
            f"Train Loss: {train_loss:.10f} | "
            f"Val Loss: {val_loss:.10f}"
        )

    if interactive_plot:
        plot_loss_curves(train_losses, val_losses)
    return train_losses, val_losses


# Plotly interactive loss curve
def plot_loss_curves(train_losses, val_losses):
    epochs = list(range(1, len(train_losses) + 1))
    trace1 = go.Scatter(x=epochs, y=train_losses, mode='lines+markers', name='Train Loss')
    trace2 = go.Scatter(x=epochs, y=val_losses, mode='lines+markers', name='Val Loss')
    layout = go.Layout(
        title='Training and Validation Loss',
        xaxis=dict(title='Epoch'),
        yaxis=dict(title='Loss', type='linear'),
        legend=dict(x=0, y=1),
        hovermode='closest',
    )
    fig = go.Figure(data=[trace1, trace2], layout=layout)
    # pyo.plot(fig, filename='loss_curve.html', auto_open=True)


# ======================================================
# 5. EVALUATION
# ======================================================
def evaluate_model(model, X_test, y_test):
    preds = model(X_test, training=False)
    mse = np.mean((preds.numpy() - y_test) ** 2)
    print(f"Test MSE: {mse:.10f}")
    return preds


# ======================================================
# 6. MAIN EXECUTION
# ======================================================
if __name__ == "__main__":
    # ---- Load dataset ----
    X, y = load_csv_dataset(CSV_FILE)
    print(f"Loaded data: X shape {X.shape}, y shape {y.shape}")
    # ---- Split train/val/test ----
    X_train, X_temp, y_train, y_temp = train_test_split(
        X, y, test_size=0.3, random_state=42
    )
    X_val, X_test, y_val, y_test = train_test_split(
        X_temp, y_temp, test_size=0.5, random_state=42
    )
    # ---- Create datasets ----
    train_ds, val_ds = create_tf_datasets(
        X_train, y_train, X_val, y_val, batch_size=32
    )
    # ---- Build model ----
    model = DenseRegressionModel(
        hidden_units1=16, hidden_units2=16, output_dim=y.shape[1]
    )
    model(np.zeros((1, X.shape[1]), dtype=np.float32))
    print(model.summary())
    # ---- Train model ----
    train_losses, val_losses = train_model(model, train_ds, val_ds, epochs=1, lr=1e-3, interactive_plot=True)
    # ---- Evaluate ----
    preds = evaluate_model(model, X_test, y_test)

    # ---- Save model as TFLite ----
    # Convert to TFLite
    converter = tf.lite.TFLiteConverter.from_keras_model(model)
    tflite_model = converter.convert()
    # Save to file
    with open("model_t6/model_t6.tflite", "wb") as f:
        f.write(tflite_model)
    print("TFLite model saved as model_t6/model_t6.tflite")
