import lightnet as ln
import torch
import numpy as np, pandas as pd
import matplotlib.pyplot as plt
import brambox as bb
import dask as da
from datasets import BramboxPathFlowDataset
import argparse, pickle
from sklearn.model_selection import train_test_split

# Settings
ln.logger.setConsoleLevel('ERROR')             # Only show error log messages
bb.logger.setConsoleLevel('ERROR')
# https://eavise.gitlab.io/lightnet/notes/02-B-engine.html

p=argparse.ArgumentParser()
p.add_argument('--num_classes',default=4,type=int)
p.add_argument('--patch_size',default=512,type=int)
p.add_argument('--patch_info_file',default='cell_info.db',type=str)
p.add_argument('--input_dir',default='inputs',type=str)
p.add_argument('--sample_p',default=1.,type=float)
p.add_argument('--conf_thresh',default=0.01,type=float)
p.add_argument('--nms_thresh',default=0.5,type=float)


args=p.parse_args()
np.random.seed(42)
num_classes=args.num_classes+1
patch_size=args.patch_size
batch_size=64
patch_info_file=args.patch_info_file
input_dir=args.input_dir
sample_p=args.sample_p
conf_thresh=args.conf_thresh
nms_thresh=args.nms_thresh
anchors=pickle.load(open('anchors.pkl','rb'))

annotation_file = 'annotations_bbox_{}.pkl'.format(patch_size)
annotations=bb.io.load('pandas',annotation_file)

if sample_p < 1.:
    annotations=annotations.sample(frac=sample_p)

annotations_dict={}
annotations_dict['train'],annotations_dict['test']=train_test_split(annotations)
annotations_dict['train'],annotations_dict['val']=train_test_split(annotations_dict['train'])

model=ln.models.Yolo(num_classes=num_classes,anchors=anchors.tolist())

loss = ln.network.loss.RegionLoss(
    num_classes=model.num_classes,
    anchors=model.anchors,
    stride=model.stride
)

transforms = ln.data.transform.Compose([ln.data.transform.RandomHSV(
    hue=1,
    saturation=2,
    value=2
)])

# Create HyperParameters
params = ln.engine.HyperParameters(
    network=model,
    input_dimension = (patch_size,patch_size),
    mini_batch_size=16,
    batch_size=batch_size,
    max_batches=80000
)

post = ln.data.transform.Compose([
    ln.data.transform.GetBoundingBoxes(
        num_classes=params.network.num_classes,
        anchors=params.network.anchors,
        conf_thresh=conf_thresh,
    ),

    ln.data.transform.NonMaxSuppression(
        nms_thresh=nms_thresh
    ),

    ln.data.transform.TensorToBrambox(
        network_size=(patch_size,patch_size),
        # class_label_map=class_label_map,
    )
])

datasets={k:BramboxPathFlowDataset(input_dir,patch_info_file, patch_size, annotations_dict[k], input_dimension=(patch_size,patch_size), class_label_map=None, identify=None, img_transform=None, anno_transform=None) for k in ['train','val','test']}
# transforms

params.loss = ln.network.loss.RegionLoss(params.network.num_classes, params.network.anchors)
params.optim = torch.optim.SGD(params.network.parameters(), lr=1e-4)
params.scheduler = ln.engine.SchedulerCompositor(
    #   batch   scheduler
        (0,     torch.optim.lr_scheduler.CosineAnnealingLR(params.optim,T_max=200))
    )

dls = {k:ln.data.DataLoader(
    datasets[k],
    batch_size = batch_size,
    collate_fn = ln.data.brambox_collate   # We want the data to be grouped as a list
    ) for k in ['train','val','test']}

params.val_loader=dls['val']

class CustomEngine(ln.engine.Engine):
    def start(self):
        """ Do whatever needs to be done before starting """
        self.params.to(self.device)  # Casting parameters to a certain device
        self.optim.zero_grad()       # Make sure to start with no gradients
        self.loss_acc = []           # Loss accumulator

    def process_batch(self, data):
        """ Forward and backward pass """
        data, target = data  # Unpack
        #print(target)
        data=data.permute(0,3,1,2).float()
        if torch.cuda.is_available():
            data=data.cuda()

        #print(data)

        output = self.network(data)
        #print(output)

        loss = self.loss(output, target)

        #print(loss)
        loss.backward()
        bbox=post(output)
        print(bbox)

        self.loss_acc.append(loss.item())

    @ln.engine.Engine.batch_end(100) # how to pass in validation dataloader
    def val_loop(self):
        with torch.no_grad():
            for i,data in enumerate(self.val_loader):
                if i > 100:
                    break
                data, target = data
                data=data.permute(0,3,1,2).float()
                if torch.cuda.is_available():
                    data=data.cuda()
                output = self.network(data)
                #print(output)
                loss = self.loss(output, target)
                print(loss)
                bbox=post(output)
                print(bbox)
                if not i:
                    bbox_final=[bbox]
                else:
                    bbox_final.append(bbox)

            detections=pd.concat(bbox_final)
            print(detections)
            print(annotations_dict['val'])
            pr=bb.stat.pr(detections, annotations_dict['val'], threshold=0.5)
            auc=bb.stat.auc(pr)
            print('VAL AUC={}'.format(auc))

    @ln.engine.Engine.batch_end(300)
    def save_model(self):
        self.params.save(f'backup-{self.batch}.state.pt')

    def train_batch(self):
        """ Weight update and logging """
        self.optim.step()
        self.optim.zero_grad()

        batch_loss = sum(self.loss_acc) / len(self.loss_acc)
        self.loss_acc = []
        self.log(f'Loss: {batch_loss}')

    def quit(self):
        if self.batch >= self.max_batches:  # Should probably save weights here
            print('Reached end of training')
            return True
        return False



# Create engine
engine = CustomEngine(
    params, dls['train'],              # Dataloader (None) is not valid
    device=torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
)

for i in range(10):
    engine()