#%%
import torch
import torchinfo
from OMI.RectifiedFlow import RectifiedFlow
from OMI.dit import DiT
from OMI.utils import RunningMean

device='cuda'
torch.set_float32_matmul_precision('high')

import os
from PIL import Image
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
from torchvision.transforms import v2

class ImageDataset(Dataset):
    def __init__(self, img_dir, transform=None):
        """
        Args:
            img_dir (string): Directory with all the images.
            transform (callable, optional): Optional transform to be applied
                on an image.
        """
        self.img_dir = img_dir
        self.image_names = os.listdir(img_dir)
        self.transform = transform
        
    def __len__(self):
        return len(self.image_names)
        
    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()
        
        # Get image path
        img_name = self.image_names[idx]
        img_path = os.path.join(self.img_dir, img_name)
        
        # Load image
        image = Image.open(img_path).convert('RGB')  # Convert to RGB
        
        # Apply transformations if any
        if self.transform:
            image = self.transform(image)
        
        return image


#%%
def main():
    model=DiT(384,8,3,3,1,8,2)
    print(torchinfo.summary(model, input_size=[(1,3,128,128),(1,1,1,1)]))

    modelT=RectifiedFlow(model,1000,device=device,lr=1e-4,emaStrength=0.999)

    #%%    
    from torch.utils.tensorboard import SummaryWriter
    import datetime
    logger=SummaryWriter(f'./omitest/{datetime.datetime.now().strftime("%Y%m%d-%H%M%S")}')
    #%%
    dataset = ImageDataset(img_dir='/mnt/h/minecraft/icons', transform=v2.Compose([
        v2.PILToTensor(),
        v2.ToDtype(torch.float32, scale=True),
        v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ]))

# Optionally, create a DataLoader to iterate over the dataset
    dataloader = DataLoader(dataset, batch_size=32, shuffle=True)
    #%%
    #Load weights
    #checkpoint = torch.load('./checkpoint.pt',map_location='cpu',weights_only=True)
    #modelT.load_state_dict(checkpoint)
    #%%
    step=-1
    while True:
        #Training
        loss=RunningMean()
        for image in dataloader:
            #image=torch.nn.functional.interpolate(image,(128,128))
            metrics=modelT.train_step(image.to('cuda'))
            loss.update(metrics['loss'])
            
            step+=1
            print(f"Step {step}, Loss: {metrics['loss']}",end='\r')
        
            if step%100==0:
                break

        #Logging
        logger.add_scalar(f'Generator Loss', loss.value, step)
        if(step % 100 == 0):
            with torch.inference_mode(), torch.amp.autocast('cuda',dtype=torch.float32):
                pred=modelT.call(steps=20,shape=(1,3,128,128))
                for i in range(1):
                    logger.add_image(f'Generated Image/{i}', pred[i].float(), step)
                    logger.add_image(f'Real Image/{i}', image[i], step)
                logger.flush()

        #Save all models
        #torch.save(modelT.state_dict(),'./checkpoint.pt')

if __name__ == '__main__':
    main()


i got noise as generated images, i think its because i normalizing images, how to fix them?