using Microsoft.ML;
using Microsoft.ML.Tokenizers;
using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using TorchSharp;

var batch = 1;
var bpe = new Bpe("vocab.json", "merges.txt", endOfWordSuffix: "</w>");
var tokenier = new Tokenizer(bpe);
var start_token = 49406;
var end_token = 49407;
var prompt = "a wild cute green cat";
var res = tokenier.Encode(prompt);
var tokens = new[] { start_token }.Concat(res.Ids.Concat(Enumerable.Repeat(0, 75 - res.Ids.Count))).Concat(new[] { end_token }).ToList();
var uncontional_tokens = new[]{start_token, end_token}.Concat(Enumerable.Repeat(0, 75)).ToList();
var tokenTensor = torch.tensor(tokens.ToArray(), dtype: torch.ScalarType.Int64, device: device);
tokenTensor = tokenTensor.repeat(batch, 1);
var unconditional_tokenTensor = torch.tensor(uncontional_tokens.ToArray(), dtype: torch.ScalarType.Int64, device: device);
unconditional_tokenTensor = unconditional_tokenTensor.repeat(batch, 1);

torchvision.io.DefaultImager = new torchvision.io.SkiaImager();
var device = TorchSharp.torch.device("cuda:0");
var clipEncoder = new ClipEncoder("clip_encoder.ckpt", device);
var img = torch.randn(batch, 4, 64, 64, dtype: torch.ScalarType.Float32, device: device);
var t = torch.full(new[]{batch, 1L}, value: batch, dtype: torch.ScalarType.Int32, device: device);
var condition = clipEncoder.Forward(tokenTensor);
var unconditional_condition = clipEncoder.Forward(unconditional_tokenTensor);

clipEncoder.Dispose();
var ddpm = new DDPM("ddim_v_sampler.ckpt", device);
var ddimSampler = new DDIMSampler(ddpm);
var ddim_steps = 50;
img = ddimSampler.Sample(img, condition, unconditional_condition, ddim_steps);
ddpm.Dispose();

var autoencoderKL = new AutoencoderKL("autoencoder_kl.ckpt", device);
var decoded_images = (torch.Tensor)autoencoderKL.Forward(img);
decoded_images = torch.clamp((decoded_images + 1.0) / 2.0, 0.0, 1.0);


for(int i = 0; i!= batch; ++i)
{
    var image = decoded_images[i];
    image = (image * 255.0).to(torch.ScalarType.Byte).cpu();
    torchvision.io.write_image(image, $"{i}.png", torchvision.ImageFormat.Png);
}