# Copyright © 2025 UChicago Argonne, LLC All right reserved
# Full license accessible at https://github.com//AdvancedPhotonSource/pty-chi/blob/main/LICENSE

from typing import Optional, Type, Union
import dataclasses
from dataclasses import field


import ptychi.api.options.base as base
import ptychi.api.options.task as task_options
import ptychi.api.enums as enums
import ptychi.api.options.ad_general as ad_general
import ptychi.forward_models as fm


@dataclasses.dataclass
class AutodiffPtychographyReconstructorOptions(ad_general.AutodiffReconstructorOptions):
    forward_model_class: Union["enums.ForwardModels", Type["fm.ForwardModel"]] = enums.ForwardModels.PLANAR_PTYCHOGRAPHY
    
    def get_reconstructor_type(self) -> enums.Reconstructors:
        return enums.Reconstructors.AD_PTYCHO


@dataclasses.dataclass
class DeepImagePriorOptions(base.Options):
    enabled: bool = False
    """Whether to use deep image prior."""
    
    model: enums.DIPModels = enums.DIPModels.UNET
    """The model used to represent the reconstruction parameter."""
    
    model_params: dict = dataclasses.field(default_factory=dict)
    """The parameters for the model's constructor."""
    
    net_input_channels: int = 32
    """The number of channels in the input to the network."""
    
    constrain_object_outside_network: bool = False
    """If True, magnitude and phase of the object generated by the model is
    further constrained so that the magnitude is between 0 and 1 outside the
    network model. Do not enable this option if your model already includes
    such constraints.
    """
    
    residual_generation: bool = False
    """If True, data generated by the DIP network is added to the initial guess
    of the reconstruction parameter, i.e., the network learns a residue. When
    using this feature, the network is usually supposed to start from a state
    that maps everything to zeros, so that the overall DIP generation is an
    identity function.
    """
    
    
@dataclasses.dataclass
class AutodiffPtychographyObjectExperimentalOptions(base.Options):
    deep_image_prior_options: Optional[DeepImagePriorOptions] = dataclasses.field(default_factory=DeepImagePriorOptions)


@dataclasses.dataclass
class AutodiffPtychographyObjectOptions(base.ObjectOptions):
    experimental: AutodiffPtychographyObjectExperimentalOptions = dataclasses.field(default_factory=AutodiffPtychographyObjectExperimentalOptions)


@dataclasses.dataclass
class AutodiffPtychographyProbeExperimentalOptions(base.Options):
    deep_image_prior_options: Optional[DeepImagePriorOptions] = dataclasses.field(default_factory=DeepImagePriorOptions)


@dataclasses.dataclass
class AutodiffPtychographyProbeOptions(base.ProbeOptions):
    experimental: AutodiffPtychographyProbeExperimentalOptions = dataclasses.field(default_factory=AutodiffPtychographyProbeExperimentalOptions)


@dataclasses.dataclass
class AutodiffPtychographyProbePositionOptions(base.ProbePositionOptions):
    pass


@dataclasses.dataclass
class AutodiffPtychographyOPRModeWeightsOptions(base.OPRModeWeightsOptions):
    pass


@dataclasses.dataclass
class AutodiffPtychographyOptions(task_options.PtychographyTaskOptions):
    reconstructor_options: AutodiffPtychographyReconstructorOptions = field(
        default_factory=AutodiffPtychographyReconstructorOptions
    )

    object_options: AutodiffPtychographyObjectOptions = field(
        default_factory=AutodiffPtychographyObjectOptions
    )

    probe_options: AutodiffPtychographyProbeOptions = field(
        default_factory=AutodiffPtychographyProbeOptions
    )

    probe_position_options: AutodiffPtychographyProbePositionOptions = field(
        default_factory=AutodiffPtychographyProbePositionOptions
    )

    opr_mode_weight_options: AutodiffPtychographyOPRModeWeightsOptions = field(
        default_factory=AutodiffPtychographyOPRModeWeightsOptions
    )
