import os

dist_params = dict(backend="nccl")
log_level = "INFO"
load_from = None
resume_from = None
cudnn_benchmark = True
custom_imports = dict(imports=["geospatial_fm"])
num_frames = 5
img_size = 224
num_workers = 16

# model
# TO BE DEFINED BY USER: model path
# pretrained_weights_path = "<path to pretrained weights>"
pretrained_weights_path="your_path/Prithvi_100M.pt"
num_layers = 6
patch_size = 16
embed_dim = 768
num_heads = 8
tubelet_size = 1
max_epochs = 50
eval_epoch_interval = 1

loss_weights_multi=[
    [1.4970547, 4.486091],
]

loss_func = dict(
    type="MultiCropBCE",
    use_sigmoid=False,
    loss_weight=1,
    class_weight=loss_weights_multi,
    avg_non_ignore=True,
)

output_embed_dim = embed_dim * num_frames


# TO BE DEFINED BY USER: Save directory
experiment="your_path/experiment_path"
project_dir="your_path/hls-foundation-os"
work_dir = os.path.join(project_dir, experiment)
save_path = work_dir

dataset_type = "GeospatialDataset"

# TO BE DEFINED BY USER: data directory
# data_root = "<path to data root>"
data_root="your_path/data_root"

splits = dict(
    # train="<path to train split>",
    # val="<path to val split>",
    # test="<path to test split>",
    train='your_path/data_root/training.txt',       #your training chip names
    val='your_path/data_root/validation.txt',       #your validation chip names
    test='your_path/data_root/test.txt',            #your test chip names
)

img_norm_cfg = dict(
    means=[
        450.19196, 745.92365, 826.1751, 2313.818, 2176.9692, 1426.2015,
        433.5821, 770.37317, 746.4662, 3024.954, 2187.981, 1327.2589,
        380.36334, 695.8719, 616.4835, 3414.738, 2118.7446, 1150.6, 332.9584,
        628.2434, 543.941, 3289.7153, 1953.0919, 1009.1733, 369.02734,
        651.11707, 672.8113, 2699.3572, 2067.684, 1184.0001,
    ],
    stds=[
       152.56088, 202.49983, 293.4363, 563.56287, 636.2972, 525.2601,
        182.19481, 233.54333, 352.48758, 669.775, 673.49335, 586.34436,
        176.17729, 232.39455, 340.19424, 671.1155, 659.76984, 552.55646,
        138.53162, 183.13971, 262.7548, 680.1936, 544.2928, 431.70184,
        151.47745, 202.88774, 299.03357, 583.2363, 605.8754, 483.53812,
    ],
)

bands = [0, 1, 2, 3, 4, 5]
tile_size = 224
orig_nsize = 224
crop_size = (tile_size, tile_size)

train_pipeline = [
    dict(type="LoadGeospatialImageFromFile", to_float32=True),
    dict(type="LoadGeospatialAnnotations", reduce_zero_label=False),
    dict(type="RandomFlip", prob=0),
    dict(type="ToTensor", keys=["img", "gt_semantic_seg"]),
    # to channels first
    dict(type="TorchPermute", keys=["img"], order=(2, 0, 1)),
    dict(type="TorchNormalize", **img_norm_cfg),
    dict(type="TorchRandomCrop", crop_size=crop_size),
    dict(
        type="Reshape",
        keys=["img"],
        new_shape=(len(bands), num_frames, tile_size, tile_size),
    ),
    dict(type="Reshape", keys=["gt_semantic_seg"], new_shape=(1, tile_size, tile_size)),
    dict(type="CastTensor", keys=["gt_semantic_seg"], new_type="torch.LongTensor"),
    dict(type="Collect", keys=["img", "gt_semantic_seg"]),
]

test_pipeline = [
    dict(type="LoadGeospatialImageFromFile", to_float32=True),
    dict(type="ToTensor", keys=["img"]),
    # to channels first
    dict(type="TorchPermute", keys=["img"], order=(2, 0, 1)),
    dict(type="TorchNormalize", **img_norm_cfg),
    dict(
        type="Reshape",
        keys=["img"],
        new_shape=(len(bands), num_frames, -1, -1),
        look_up=dict({"2": 1, "3": 2}),
    ),
    dict(type="CastTensor", keys=["img"], new_type="torch.FloatTensor"),
    dict(
        type="CollectTestList",
        keys=["img"],
        meta_keys=[
            "img_info",
            "seg_fields",
            "img_prefix",
            "seg_prefix",
            "filename",
            "ori_filename",
            "img",
            "img_shape",
            "ori_shape",
            "pad_shape",
            "scale_factor",
            "img_norm_cfg",
        ],
    ),
]

CLASSES=[(
    "Other",
    'Rice',
)]

dataset = "GeospatialDataset"
data = dict(
    samples_per_gpu=6,      
    workers_per_gpu=8,
    train=dict(
        type=dataset,
        CLASSES=CLASSES,
        data_root=data_root,
        img_dir="training_chips",   #training chips dir
        ann_dir="training_chips",   
        pipeline=train_pipeline,
        img_suffix="_merged.tif",
        seg_map_suffix=".mask.tif",
        split=splits["train"],
    ),
    val=dict(
        type=dataset,
        CLASSES=CLASSES,
        data_root=data_root,
        img_dir="rice_validation_chips",    #validation chips dir
        ann_dir="rice_validation_chips",
        pipeline=test_pipeline,
        img_suffix="_merged.tif",
        seg_map_suffix=".mask.tif",
        split=splits["val"],
    ),
    test=dict(
        type=dataset,
        CLASSES=CLASSES,
        data_root=data_root,
        img_dir="rice_test_chips",          #test chips dir
        ann_dir="rice_test_chips",
        pipeline=test_pipeline,
        img_suffix="_merged.tif",
        seg_map_suffix=".mask.tif",
        split=splits["test"],
    ),
)

optimizer = dict(type="Adam", lr=5e-05, betas=(0.9, 0.999), weight_decay=0.05)
optimizer_config = dict(grad_clip=None)
lr_config = dict(
    policy="poly",
    warmup="linear",
    warmup_iters=1500,
    warmup_ratio=1e-06,
    power=1.0,
    min_lr=0.0,
    by_epoch=False,
)
log_config = dict(
    interval=10, hooks=[dict(type="TextLoggerHook"), dict(type="TensorboardLoggerHook")]
)

checkpoint_config = dict(by_epoch=True, interval=100, out_dir=save_path)

evaluation = dict(
    interval=eval_epoch_interval,
    metric="mIoU",
    pre_eval=True,
    save_best="mIoU",
    by_epoch=True,
)
reduce_train_set = dict(reduce_train_set=False)
reduce_factor = dict(reduce_factor=1)
runner = dict(type="EpochBasedRunner", max_epochs=max_epochs)
workflow = [("train", 1)]
norm_cfg = dict(type="BN", requires_grad=True)

model = dict(
    type="TemporalEncoderDecoder",
    frozen_backbone=False,
    backbone=dict(
        type="TemporalViTEncoder",
        pretrained=pretrained_weights_path,
        img_size=img_size,
        patch_size=patch_size,
        num_frames=num_frames,
        tubelet_size=1,
        in_chans=len(bands),
        embed_dim=embed_dim,
        depth=6,
        num_heads=num_heads,
        mlp_ratio=4.0,
        norm_pix_loss=False,
    ),
    neck=dict(
        type="ConvTransformerTokensToEmbeddingNeck",
        embed_dim=embed_dim * num_frames,
        output_embed_dim=output_embed_dim,
        drop_cls_token=True,
        Hp=14,
        Wp=14,
        num_frames=num_frames,
        in_channels=len(bands)
    ),
    decode_head=dict(
        num_classes=2,
        in_channels=output_embed_dim+num_frames*len(bands),
        type="CropHead",
        in_index=-1,
        channels=256,
        num_convs=1,
        concat_input=False,
        dropout_ratio=0.1,
        norm_cfg=dict(type="BN", requires_grad=True),
        align_corners=False,
        loss_decode=loss_func,
        ignore_index=3,
    ),
    auxiliary_head=dict(
        num_classes=len(CLASSES)*2,
        in_channels=output_embed_dim+num_frames*len(bands),
        type="CropHead",
        in_index=-1,
        channels=256,
        num_convs=2,
        concat_input=False,
        dropout_ratio=0.1,
        norm_cfg=dict(type="BN", requires_grad=True),
        align_corners=False,
        loss_decode=loss_func,
        ignore_index=3,
    ),
    train_cfg=dict(),
    test_cfg=dict(
        mode="slide",
        stride=(int(tile_size / 2), int(tile_size / 2)),
        crop_size=(tile_size, tile_size),
    ),
)
auto_resume = False