from tfx.types.component_spec import ChannelParameter, ComponentSpec, ExecutionParameter
from tfx.types import standard_artifacts

class MyComponentSpec(ComponentSpec):
    PARAMETERS = {
        "num_samples": ExecutionParameter(type=int),
        "config": ExecutionParameter(type=str),
    }
    INPUTS = {
        "examples": ChannelParameter(type=standard_artifacts.Examples),
    }
    OUTPUTS = {
        "model": ChannelParameter(type=standard_artifacts.Model),
    }

from tfx import types
from tfx.dsl.components.base import base_component, executor_spec
from tfx.types import standard_artifacts
from my_component import executor

class MyComponent(base_component.BaseComponent):
    SPEC_CLASS = MyComponentSpec
    EXECUTOR_SPEC = executor_spec.ExecutorClassSpec(executor.Executor)

    def __init__(self,
                 num_samples: int,
                 config: Dict[str, str],
                 examples: types.Channel = None):
    output_model = types.Channel(type=standard_artifacts.Model)

    spec = MyComponentSpec(num_samples=num_samples,
                           config=json.dumps(config),
                           examples=examples,
                           model=output_model)
    super().__init__(spec=spec)

from tfx.dsl.components.base import base_executor
from tfx.types import artifact_utils

class Executor(base_executor.Executor):
    def Do(
        self,
        input_dict: Dict[str, List[Artifact]],
        output_dict: Dict[str, List[Artifact]],
        exec_properties: Dict[str, Any],
    ):
        num_samples = exec_properties.get("num_samples")
        config = json.loads(exec_properties.get("config"))
        examples_artifact = artifact_utils.get_single_instance(input_dict["examples"])
        model_artifact = artifact_utils.get_single_instance(output_dict["model"])

        # 아래처럼 경로에서 데이터를 불러와야 합니다.
        # load_examples() 는 가상의 함수로, 아티팩트 경로에서 데이터를 불러옵니다.
        examples = load_examples(examples_artifact.uri)

        # 로직을 작성합니다.

        # 아티팩트는 uri 속성에 자신이 바라보는 경로가 지정되어 있습니다.
        # write_model()는 가상의 함수로, examples를 입력으로 받아 model을 학습하고나서
        # 최상의 checkpoint를 model_artifact에 기록합니다.
        write_model(model_artifact.uri, examples)

# Component Interface
from tfx.dsl.components.base import base_beam_component

class MyBeamComponent(base_beam_component.BaseBeamComponent):
    pass

# Executor
from tfx.dsl.components.base import base_beam_executor

class Executor(base_beam_executor.BaseBeamExecutor):
    def Do(
        self,
        input_dict: Dict[str, List[Artifact]],
        output_dict: Dict[str, List[Artifact]],
        exec_properties: Dict[str, Any],
    ):

        with self._make_beam_pipeline() as pipeline:
            # Pipeline 작성

class _DatasetsExampleGenSpec(ComponentSpec):
    PARAMETERS = {
        "transform_fn": ExecutionParameter(type=str, optional=True),
        "module_file": ExecutionParameter(type=str, optional=True),
        "input_config": ExecutionParameter(type=example_gen_pb2.Input),
        "output_config": ExecutionParameter(type=example_gen_pb2.Output),
        "custom_config": ExecutionParameter(type=str, optional=True),
        "range_config": ExecutionParameter(type=range_config_pb2.RangeConfig, optional=True),
        "builder_config": ExecutionParameter(type=str, optional=True),
    }
    INPUTS = {}
    OUTPUTS = {
        "examples": ChannelParameter(type=standard_artifacts.Examples),
    }
    TYPE_ANNOTATION = Process

class DatasetsExampleGen(base_beam_component.BaseBeamComponent):
    """Dataset Registry 로부터 데이터를 불러와 Example을 생성합니다."""

    SPEC_CLASS = _DatasetsExampleGenSpec
    EXECUTOR_SPEC = executor_spec.BeamExecutorSpec(executor.Executor)
    DRIVER_CLASS = driver.DatasetsDriver

    def __init__(
        self,
        name: Optional[str] = None,
        transform_fn: Optional[Union[str, data_types.RuntimeParameter]] = None,
        module_file: Optional[Union[str, data_types.RuntimeParameter]] = None,
        input_config: Optional[Union[example_gen_pb2.Input, data_types.RuntimeParameter]] = None,
        output_config: Optional[Union[example_gen_pb2.Output, data_types.RuntimeParameter]] = None,
        custom_config: Optional[Union[Dict[str, Any], data_types.RuntimeParameter]] = None,
        range_config: Optional[Union[range_config_pb2.RangeConfig, data_types.RuntimeParameter]] = None,
        builder_config: Optional[Union[Dict[str, Any], data_types.RuntimeParameter]] = None,
    ):
        if name and input_config:
            raise RuntimeError("Exactly one of name and input_config should be set.")

        input_config = input_config or utils.make_default_input_config(name)
        output_config = output_config or utils.make_default_output_config(input_config)
        example_artifacts = types.Channel(type=standard_artifacts.Examples)

        spec = _DatasetsExampleGenSpec(
            transform_fn=transform_fn,
            module_file=module_file,
            input_config=input_config,
            output_config=output_config,
            range_config=range_config,
            custom_config=json.dumps(custom_config),
            examples=example_artifacts,
            builder_config=json.dumps(builder_config),
        )

        super().__init__(spec=spec)

        if udf_utils.should_package_user_modules():
            udf_utils.add_user_module_dependency(
                self,
                standard_component_specs.MODULE_FILE_KEY,
                standard_component_specs.MODULE_PATH_KEY,
            )

class _TransformToSerialized(beam.DoFn):
    def process(self, example: Any):
        yield from self.transform_fn(example)

class _DatasetsToExample(beam.PTransform):
    def expand(self, pipeline: beam.Pipeline) -> beam.PCollection[tf.train.Example]:
        return (
            pipeline
            | "ReadFromDatasets" >> datasets.ReadFromDatasets(self.name, self.version)
            | "ParseExample" >> beam.ParDo(_TransformToSerialized(self.name, self.transform_fn))
            | "ToTFExample" >> beam.Map(tf.train.Example.FromString)
        )

class Executor(BaseExampleGenExecutor):
    def GetInputSourceToExamplePTransform(self) -> beam.PTransform:
        return functools.partial(_DatasetsToExample, transform_fn=self.transform_fn)

class DatasetsDriver(driver.Driver):
    def get_input_processor(
        self,
        splits: Iterable[example_gen_pb2.Input.Split],
        range_config: Optional[range_config_pb2.RangeConfig] = None,
        input_base_uri: Optional[str] = None,
    ) -> input_processor.InputProcessor:
        return DatasetsInputProcessor(splits, range_config)

class DatasetsInputProcessor(input_processor.InputProcessor):
    def get_latest_span(self) -> int:
        versions = datasets.ls(pattern)
        # 불러온 데이터셋 목록에서 가장 최신의 버전을 찾습니다.
        return datasets_utils.get_lastest_span(versions)

    def get_latest_version(self, span: int) -> Optional[int]:
        versions = datasets.ls(pattern)
        # 불러온 데이터셋 목록에서 가장 최신의 버전을 찾습니다.
        return datasets_utils.get_latest_version(versions, span)

input_config = example_gen_pb2.Input(
    example_gen_pb2.Input.Split(name="train", pattern="klue_sts/v{SPAN}_{VERSION}-train"),
    example_gen_pb2.Input.Split(name="eval", pattern="klue_sts/v{SPAN}_{VERSION}-test"),
)
example_gen = DatasetsExampleGen(input_config=input_config)

def preprocessing_fn(inputs: Dict[str, Any], custom_config: Dict[str, Any]) -> Dict[str, Any]:
    tokenizer = tft.make_and_track_object(lambda: Tokenizer(...)

class Executor(base_executor.BaseExecutor):
    def Do(
        self,
        input_dict: Dict[str, List[Artifact]],
        output_dict: Dict[str, List[Artifact]],
        exec_properties: Dict[str, Any],
    ):
        model = artifact_utils.get_single_instance(input_dict.get("model"))
        model_name = exec_properties.get("model_name")
        model_version = exec_properties.get("model_version")

        uploaded_model = artifact_utils.get_single_instance(output_dict.get("uploaded_model"))
            
        # 모델을 불러와 Neuron Compile 진행
        base_model = tf.keras.models.load_model(model.uri, compile=False)
        neuron_model = tfn.trace(base_model, _generate_inputs(model.uri))
        tf.saved_model.save(neuron_model, uploaded_model.uri, signatures=_signature_fn(model.uri))

        # Model Registry에 등록
        models.register(uploaded_model.uri, model_name, model_version)