import numpy as np
import struct
import tensorflow as tf
import json
import pickle
#from google.colab import files

def get_timestep_embedding(timestep, dim=320, max_period=10000):
    half = dim // 2
    freqs = np.exp(-np.log(max_period) * np.arange(half) / half)
    args = timestep * freqs
    embedding = np.concatenate([np.cos(args), np.sin(args)])
    return embedding.astype(np.float32)

def process_embeddings(model_path, num_steps=20):
    interpreter = tf.lite.Interpreter(model_path=model_path)
    interpreter.allocate_tensors()
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()
    print("Aswin::Type "+ str(type(output_details)))
    print("Aswin::Length "+ str(len(output_details)))
    print(output_details)

    timesteps = [999, 949, 899, 849, 799, 749, 699, 649, 599, 549, 500, 450, 400, 350, 300, 250, 200, 150, 100, 50] 
    #timesteps = list(range(1, 1000, 1000 // num_steps))
    embeddings = []
    dict1={}
    for t in timesteps:
        embedding = get_timestep_embedding(t)
        interpreter.set_tensor(input_details[0]['index'], embedding.reshape(1, 320))
        interpreter.invoke()
        output = interpreter.get_tensor(output_details[0]['index'])#.reshape(-1)
        embeddings.append(output)
        ind = timesteps.index(t)
        dict1[str(ind)]=output
    return timesteps, embeddings, dict1

def save_json(timesteps, embeddings, output_path):
    data = {
        "timestep_sequences": {
            str(len(timesteps)): timesteps
        },
        "embeddings": {
            str(len(timesteps)): [
                {
                    "step": i,
                    "embedding": emb.tolist(),
                    "embedding_size": len(emb),
                    "min": float(np.min(emb)),
                    "max": float(np.max(emb)),
                    "avg": float(np.mean(emb))
                }
                for i, emb in enumerate(embeddings)
            ]
        }
    }

    with open(output_path, 'w') as f:
        json.dump(data, f, indent=2)
    #files.download(output_path)

def save_pickle(timesteps, embeddings, dict1, output_path):
    #data = {
        #"timesteps": timesteps,
        #"embeddings": embeddings,
     #   dict1
        #"metadata": {
        #    "num_steps": len(timesteps),
        #    "embedding_size": embeddings[0].shape[0] if embeddings else 0
        #}
    #}
    with open(output_path, 'wb') as f:
        pickle.dump(dict1, f)
    #files.download(output_path)

model_path = '/local/mnt/workspace/aswib/timestep_expt/head_cut_models_SD/sdv15_time_emb_head_LinSiluLin_fp32.tflite'
timesteps, embeddings, dict1 = process_embeddings(model_path)

save_json(timesteps, embeddings, '/local/mnt/workspace/aswib/timestep_expt/timestep_embeddings.json')
save_pickle(timesteps, embeddings, dict1, '/local/mnt/workspace/aswib/timestep_expt/timestep_embeddings.pkl')