from climatenet import *
import rasterio



print("main.py started")


def load_raster(file_path: str):
    with rasterio.open(file_path) as src:
        data = src.read()

        tensor = torch.from_numpy(data.astype(np.float32))
        return tensor
    

def save_raster_like(tensor: torch.Tensor, ref_path: str, file_path: str):
    data = tensor.detach().cpu().numpy()

    with rasterio.open(ref_path) as example:
        meta = example.meta.copy()

    meta.update({
        "driver": "GTiff",
        "height": data.shape[1],
        "width":  data.shape[2],
        "count":  data.shape[0],
        "dtype":  data.dtype
    })

    with rasterio.open(file_path, "w", **meta) as dst:
        dst.write(data)



if __name__ == "__main__":
    torch.manual_seed(46)


    building_height = load_raster("data/INPUT/building_height.tif")
    tree_height = load_raster("data/INPUT/tree_height.tif")
    surface_height = load_raster("data/INPUT/zt.tif")
    pavement_type = load_raster("data/INPUT/pavement_type.tif")
    vegetation_type = load_raster("data/INPUT/vegetation_type.tif")
    water_type = load_raster("data/INPUT/water_type.tif")

    output_pet = load_raster("data/OUTPUT/bio_pet_xy_av_14h.tif")
    output_temp = load_raster("data/OUTPUT/ta_av_h001_1.0m_14h.tif")



    building_height = F.pad(building_height, (3, 3, 3, 3), mode='constant', value=0).unsqueeze(0)
    tree_height = F.pad(tree_height, (3, 3, 3, 3), mode='constant', value=0).unsqueeze(0)
    surface_height = F.pad(surface_height, (3, 3, 3, 3), mode='constant', value=0).unsqueeze(0)
    pavement_type = F.pad(pavement_type, (3, 3, 3, 3), mode='constant', value=0).unsqueeze(0)
    vegetation_type = F.pad(vegetation_type, (3, 3, 3, 3), mode='constant', value=0).unsqueeze(0)
    water_type = F.pad(water_type, (3, 3, 3, 3), mode='constant', value=0).unsqueeze(0)

    building_height_coarse = F.interpolate(building_height, size=(64, 64), mode='bilinear', align_corners=False)
    tree_height_coarse = F.interpolate(tree_height, size=(64, 64), mode='bilinear', align_corners=False)
    surface_height_coarse = F.interpolate(surface_height, size=(64, 64), mode='bilinear', align_corners=False)
    pavement_type_coarse = F.interpolate(pavement_type, size=(64, 64), mode='nearest')
    vegetation_type_coarse = F.interpolate(vegetation_type, size=(64, 64), mode='nearest')
    water_type_coarse = F.interpolate(water_type, size=(64, 64), mode='nearest')


    print(building_height.shape)
    print(tree_height.shape)
    print(surface_height.shape)
    print(pavement_type.shape)
    print(vegetation_type.shape)
    print(water_type.shape)
    print(output_pet.shape)
    print(output_temp.shape)


    input_data_fine = torch.cat([building_height, tree_height, surface_height, pavement_type, vegetation_type, water_type], dim=1)
    print(f"total dim fine {input_data_fine.shape}")


    input_data_coarse = torch.cat([building_height_coarse, tree_height_coarse, surface_height_coarse, pavement_type_coarse, vegetation_type_coarse, water_type_coarse], dim=1)
    print(f"total dim coarse {input_data_coarse.shape}")



    model = UrbanClimateUNet();
    model.eval()

    out = model.forward(input_data_fine, input_data_coarse)

    output_temp_pred = out.t_air_mean.squeeze(0)
    output_pet_pred = out.pet_mean.squeeze(0)

    print(f"PET pred shape: {output_pet_pred.shape}");
    print(f"PET pred shape: {output_temp_pred.shape}");

    save_raster_like(output_temp_pred, "data/OUTPUT/ta_av_h001_1.0m_14h.tif", "data/PRED/temp_coarse.tif")
    save_raster_like(output_pet_pred, "data/OUTPUT/bio_pet_xy_av_14h.tif", "data/PRED/pet.tif")