from database.database import db

class Training(db.Model):
    __tablename__ = 'training'
    
    id = db.Column(db.Integer, primary_key=True, autoincrement=True, unique=True)
    exp_name = db.Column(db.String(255))
    max_epoch = db.Column(db.Integer)
    depth = db.Column(db.Float)
    width = db.Column(db.Float)
    activation = db.Column(db.String(255))
    warmup_epochs = db.Column(db.Integer)
    warmup_lr = db.Column(db.Float)
    basic_lr_per_img = db.Column(db.Float)
    scheduler = db.Column(db.String(255))
    no_aug_epochs = db.Column(db.Integer)
    min_lr_ratio = db.Column(db.Float)
    ema = db.Column(db.Boolean)
    weight_decay = db.Column(db.Float)
    momentum = db.Column(db.Float)
    input_size = db.Column(db.JSON)
    print_interval = db.Column(db.Integer)
    eval_interval = db.Column(db.Integer)
    save_history_ckpt = db.Column(db.Boolean)
    test_size = db.Column(db.JSON)
    test_conf = db.Column(db.Float)
    nms_thre = db.Column(db.Float)
    multiscale_range = db.Column(db.Integer)
    enable_mixup = db.Column(db.Boolean)
    mosaic_prob = db.Column(db.Float)
    mixup_prob = db.Column(db.Float)
    hsv_prob = db.Column(db.Float)
    flip_prob = db.Column(db.Float)
    degrees = db.Column(db.Float)
    mosaic_scale = db.Column(db.JSON)
    mixup_scale = db.Column(db.JSON)
    translate = db.Column(db.Float)
    shear = db.Column(db.Float)
    training_name = db.Column(db.String(255))
    project_details_id = db.Column(db.Integer, nullable=False)
    seed = db.Column(db.Integer)
    train = db.Column(db.Integer)
    valid = db.Column(db.Integer)
    test = db.Column(db.Integer)
    selected_model = db.Column(db.String(255))
    transfer_learning = db.Column(db.String(255))
    model_upload = db.Column(db.LargeBinary)
    
    def to_dict(self):
        return {
            'id': self.id,
            'exp_name': self.exp_name,
            'max_epoch': self.max_epoch,
            'depth': self.depth,
            'width': self.width,
            'activation': self.activation,
            'warmup_epochs': self.warmup_epochs,
            'warmup_lr': self.warmup_lr,
            'basic_lr_per_img': self.basic_lr_per_img,
            'scheduler': self.scheduler,
            'no_aug_epochs': self.no_aug_epochs,
            'min_lr_ratio': self.min_lr_ratio,
            'ema': self.ema,
            'weight_decay': self.weight_decay,
            'momentum': self.momentum,
            'input_size': self.input_size,
            'print_interval': self.print_interval,
            'eval_interval': self.eval_interval,
            'save_history_ckpt': self.save_history_ckpt,
            'test_size': self.test_size,
            'test_conf': self.test_conf,
            'nms_thre': self.nms_thre,
            'multiscale_range': self.multiscale_range,
            'enable_mixup': self.enable_mixup,
            'mosaic_prob': self.mosaic_prob,
            'mixup_prob': self.mixup_prob,
            'hsv_prob': self.hsv_prob,
            'flip_prob': self.flip_prob,
            'degrees': self.degrees,
            'mosaic_scale': self.mosaic_scale,
            'mixup_scale': self.mixup_scale,
            'translate': self.translate,
            'shear': self.shear,
            'training_name': self.training_name,
            'project_details_id': self.project_details_id,
            'seed': self.seed,
            'train': self.train,
            'valid': self.valid,
            'test': self.test,
            'selected_model': self.selected_model,
            'transfer_learning': self.transfer_learning
        }