from google.cloud import storage


def make_bucket(project_name, bucket_name, region):
    client = storage.Client(project_name)
    bucket = storage.Bucket(client)
    bucket.name = bucket_name
    if not bucket.exists():
        client.create_bucket(bucket, location=region)

import logging

import apache_beam as beam
from apache_beam.ml.inference import RunInference
from apache_beam.ml.inference.sklearn_inference import ModelFileType, SklearnModelHandlerNumpy
from apache_beam.options.pipeline_options import PipelineOptions
from google.cloud import storage
from sklearn.datasets import load_iris


def upload_model_to_gcs(local_model_path, gcs_model_path, project_name, bucket_name):
    client = storage.Client(project_name)
    bucket = storage.Bucket(client)
    bucket.name = bucket_name
    blob = bucket.blob(gcs_model_path)
    blob.upload_from_filename(local_model_path)


if __name__ == "__main__":
    # パイプラインオプションの設定
    pipeline_options = PipelineOptions()
    options_dict = pipeline_options.display_data()

    # Irisデータの準備
    data = load_iris()
    numpy_data = data.data

    # モデルのアップロード
    upload_model_to_gcs(
        local_model_path="./models/sklearn_models/SVC_iris.pkl2",
        gcs_model_path="models/sklearn_models/SVC_iris.pkl2",
        project_name=options_dict["project"],
        bucket_name=options_dict["bucket_name"]
    )

    # ハンドラーの設定
    model_uri = "gs://{}/models/sklearn_models/SVC_iris.pkl2".format(options_dict["bucket_name"])
    model_file_type = ModelFileType.JOBLIB
    model_handler = SklearnModelHandlerNumpy(model_uri=model_uri, model_file_type=model_file_type)

    # パイプライン実行
    logging.getLogger().setLevel(logging.INFO)
    with beam.Pipeline(options=pipeline_options) as p:
        input = p | "read" >> beam.Create(numpy_data)

        prediction = (
            input
            | RunInference(model_handler)
            | beam.io.WriteToText(options_dict["output_executable_path"], shard_name_template="")
        )

import logging

import apache_beam as beam
from apache_beam.ml.inference.base import RunInference
from apache_beam.ml.inference.pytorch_inference import PytorchModelHandlerTensor
from apache_beam.options.pipeline_options import PipelineOptions
from google.cloud import storage
from torchvision import datasets, transforms

from pytorch_MNIST import MNIST_Model


def upload_model_to_gcs(local_model_path, gcs_model_path, project_name, bucket_name):
    client = storage.Client(project_name)
    bucket = storage.Bucket(client)
    bucket.name = bucket_name
    blob = bucket.blob(gcs_model_path)
    blob.upload_from_filename(local_model_path)


if __name__ == "__main__":
    # パイプラインオプションの設定
    pipeline_options = PipelineOptions()
    options_dict = pipeline_options.display_data()

    # データセットの準備
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307), (0.3081), inplace=True)
    ])
    test_dataset = datasets.MNIST(
        root="./data/",
        train=False,
        transform=transform,
        download=True
    )

    # モデルのアップロード
    upload_model_to_gcs(
        local_model_path="./models/pytorch_models/mnist_epoch_10.pth",
        gcs_model_path="models/pytorch_models/mnist_epoch_10.pth",
        project_name=options_dict["project"],
        bucket_name=options_dict["bucket_name"]
    )

    # ハンドラーの設定
    model_handler = PytorchModelHandlerTensor(
        state_dict_path="gs://{}/models/pytorch_models/mnist_epoch_10.pth".format(options_dict["bucket_name"]),
        model_class=MNIST_Model,
        model_params={},
        device="GPU"
    )

    # パイプライン実行
    logging.getLogger().setLevel(logging.INFO)
    with beam.Pipeline(options=pipeline_options) as p:
        data = p | "read" >> beam.Create(test_dataset)
        test = (
            data
            | "extract image" >> beam.Map(lambda x: x[0])
            | "inference" >> RunInference(model_handler)
            | beam.io.WriteToText(options_dict["output_executable_path"], shard_name_template="")
        )

from torch import nn


class MNIST_Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.feature = nn.Sequential(
            nn.Conv2d(1, 3, 3, padding=1),
            nn.MaxPool2d(2, 2),
            nn.ReLU(True),
            nn.Conv2d(3, 3, 3, padding=1),
            nn.MaxPool2d(2, 2),
            nn.ReLU(True)
        )
        self.classifier = nn.Sequential(
            nn.Linear(147, 128),
            nn.ReLU(True),
            nn.Linear(128, 10)
        )

    def forward(self, x):
        x = self.feature(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x

import json
import logging

import apache_beam as beam
from apache_beam.ml.inference import RunInference
from apache_beam.ml.inference.sklearn_inference import ModelFileType, SklearnModelHandlerNumpy
from apache_beam.options.pipeline_options import PipelineOptions, StandardOptions
from google.cloud import storage


def upload_model_to_gcs(local_model_path, gcs_model_path, project_name, bucket_name):
    client = storage.Client(project_name)
    bucket = storage.Bucket(client)
    bucket.name = bucket_name
    blob = bucket.blob(gcs_model_path)
    blob.upload_from_filename(local_model_path)


if __name__ == "__main__":
    # パイプラインオプションの設定
    options = PipelineOptions()
    options_dict = options.display_data()
    options.view_as(StandardOptions).runner = "DataflowRunner"
    options.view_as(StandardOptions).streaming = True

    # モデルのアップロード
    upload_model_to_gcs(
        local_model_path="./models/sklearn_models/SVC_iris.pkl2",
        gcs_model_path="models/sklearn_models/SVC_iris.pkl2",
        project_name=options_dict["project"],
        bucket_name=options_dict["bucket_name"]
    )

    # ハンドラーの設定
    model_uri = "gs://{}/models/sklearn_models/SVC_iris.pkl2".format(options_dict["bucket_name"])
    model_file_type = ModelFileType.JOBLIB
    model_handler = SklearnModelHandlerNumpy(model_uri=model_uri, model_file_type=model_file_type)

    topic = "projects/{}/topics/{}".format(options_dict["project"], options_dict["topic_name"])

    # パイプライン実行
    logging.getLogger().setLevel(logging.INFO)
    with beam.Pipeline(options=options) as p:
        raw_data = (
            p
            | "ReadFromPub/Sub" >> beam.io.ReadFromPubSub(topic)
            | "Decode" >> beam.Map(lambda x: x.decode())
            | "StrToDict" >> beam.Map(json.loads)
        )

        inference = (
            raw_data
            | "ExtractFeature" >> beam.Map(lambda x: x["feature"])
            | "RunInference" >> RunInference(model_handler)
        )

        write2bq = (
            inference
            | "ConvertToBigQueryFormat" >> beam.Map(lambda x: {
                "input": {
                    "sepal_length": x[0][0],
                    "sepal_width": x[0][1],
                    "petal_length": x[0][2],
                    "petal_width": x[0][3]
                },
                "predict": x[1].item()
            })
            | "WriteToBigQuery" >> beam.io.WriteToBigQuery(table=options_dict["table_name"], dataset=options_dict["dataset_name"])
        )

import argparse
import json
import time

from google.cloud import pubsub
from sklearn.datasets import load_iris

parser = argparse.ArgumentParser()
parser.add_argument("--project", required=True)
parser.add_argument("--topic_name", required=True)

args = parser.parse_args()

if __name__ == "__main__":
    data = load_iris()
    feature = data.data
    target = data.target

    publisher = pubsub.PublisherClient()
    topic_path = publisher.topic_path(args.project, args.topic_name)

    for i, (f, t) in enumerate(zip(feature, target)):
        f_t_dict = {"id": i, "feature": f.tolist(), "target": t.item()}
        message = json.dumps(f_t_dict)
        print(message)
        b_message = message.encode()
        publisher.publish(topic_path, b_message)
        time.sleep(10)