from argparse import ArgumentParser, Namespace

from torch.functional import norm

from ..base_model.base_model import BaseModel
from .losses import MultiScale


class FlowNetBase(BaseModel):
    def __init__(self,
                 args: Namespace):
        super(FlowNetBase, self).__init__(
            args=args,
            loss_fn=MultiScale(
                startScale=args.loss_start_scale,
                numScales=args.loss_num_scales,
                l_weight=args.loss_base_weight,
                norm=args.loss_norm
            ),
            output_stride=64)

    @staticmethod
    def add_model_specific_args(parent_parser=None):
        parent_parser = BaseModel.add_model_specific_args(parent_parser)
        parser = ArgumentParser(parents=[parent_parser], add_help=False)
        parser.add_argument('--div_flow', type=float, default=20.0)
        parser.add_argument('--input_channels', type=int, default=6)
        parser.add_argument('--batch_norm', action='store_true')
        parser.add_argument('--loss_start_scale', type=float, default=4)
        parser.add_argument('--loss_num_scales', type=int, default=5)
        parser.add_argument('--loss_base_weight', type=float, default=0.32)
        parser.add_argument('--loss_norm', type=str, default='L2')
        return parser
