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 moved to TrainingSize table
    print_interval = db.Column(db.Integer)
    eval_interval = db.Column(db.Integer)
    save_history_ckpt = db.Column(db.Boolean)
    # test_size moved to TrainingSize table
    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 moved to TrainingSize table
    # mixup_scale moved to TrainingSize table
    translate = db.Column(db.Float)
    shear = db.Column(db.Float)
    training_name = db.Column(db.String(255))
    project_details_id = db.Column(db.Integer, db.ForeignKey('training_project_details.id', ondelete='CASCADE'), 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)
    
    # Relationship to size configurations (3NF)
    size_configs = db.relationship('TrainingSize', backref='training', lazy=True, cascade='all, delete-orphan')
    
    def to_dict(self, include_sizes=True):
        result = {
            '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,
            'print_interval': self.print_interval,
            'eval_interval': self.eval_interval,
            'save_history_ckpt': self.save_history_ckpt,
            '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,
            '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
        }
        
        # Include size arrays for backwards compatibility
        if include_sizes:
            from models.TrainingSize import TrainingSize
            
            def get_size_array(size_type):
                sizes = TrainingSize.query.filter_by(
                    training_id=self.id, 
                    size_type=size_type
                ).order_by(TrainingSize.value_order).all()
                return [s.value for s in sizes] if sizes else None
            
            result['input_size'] = get_size_array('input_size')
            result['test_size'] = get_size_array('test_size')
            result['mosaic_scale'] = get_size_array('mosaic_scale')
            result['mixup_scale'] = get_size_array('mixup_scale')
        
        return result