AWSTemplateFormatVersion: "2010-09-09"
Description: Private development environment (remote to QPU region) - Sagemaker notebooks will run here and connect with the Braket VPC via inter-region VPC peering.
Parameters:
    VPCName:
      Description: Local Region VPC, located in a non-Braket Region. 
      Type: String
      Default: "Private VPC with a SageMaker development environment."

    BraketNotebookInstance:
      Description: Private Sagemaker notebook instance that will run the local Braket development environment.
      Type: String
      Default: 'amazon-braket-private'

Mappings:
  LocalSubnetConfig:
    VPC:
      CIDR: "10.20.0.0/16"
    Private1:
      CIDR: "10.20.1.0/24"
    Private2:
      CIDR: "10.20.2.0/24"

  # This mapping accounts for the scenario when certain AZs
  # are not available to use (this differs on a per account
  # per customer basis). E.g., if the 'b' AZ is not available
  # in a specific region in one's account then updating the
  # list contained in the mapping below here will allow a
  # different AZ to be chosen. 
  LocalAZRegions:
    ap-northeast-1:
      AZs: ["a", "b"]
    ap-northeast-2:
      AZs: ["a", "b"]
    ap-south-1:
      AZs: ["a", "b"]
    ap-southeast-1:
      AZs: ["a", "b"]
    ap-southeast-2:
      AZs: ["a", "b"]
    ca-central-1:
      AZs: ["a", "b"]
    eu-central-1:
      AZs: ["a", "b"]
    eu-west-1:
      AZs: ["a", "b"]
    eu-west-2:
      AZs: ["a", "b"]
    eu-west-3:
      AZs: ["a", "b"]
    sa-east-1:
      AZs: ["a", "b"]
    us-east-1:
      AZs: ["a", "b"]
    us-east-2:
      AZs: ["a", "b"]
    us-west-1:
      AZs: ["a", "b"]
    us-west-2:
      AZs: ["a", "b"]

Resources:

  VPC:
    Type: "AWS::EC2::VPC"
    Properties:
      EnableDnsSupport: "true"
      EnableDnsHostnames: "true"
      CidrBlock:
        Fn::FindInMap:
          - "LocalSubnetConfig"
          - "VPC"
          - "CIDR"
      Tags:
        - Key: Name
          Value: Braket
  
  PrivateSubnet1:
    Type: "AWS::EC2::Subnet"
    Properties:
      VpcId:
        Ref: "VPC"
      AvailabilityZone:
        Fn::Sub:
          - "${AWS::Region}${AZ}"
          - AZ: !Select [ 0, !FindInMap [ "LocalAZRegions", !Ref "AWS::Region", "AZs" ] ]
      CidrBlock:
        Fn::FindInMap:
          - "LocalSubnetConfig"
          - "Private1"
          - "CIDR"

  PrivateSubnet2:
    Type: "AWS::EC2::Subnet"
    Properties:
      VpcId:
        Ref: "VPC"
      AvailabilityZone:
        Fn::Sub:
          - "${AWS::Region}${AZ}"
          - AZ: !Select [ 1, !FindInMap [ "LocalAZRegions", !Ref "AWS::Region", "AZs" ] ]
      CidrBlock:
        Fn::FindInMap:
          - "LocalSubnetConfig"
          - "Private2"
          - "CIDR"

  RouteTable:
    Type: "AWS::EC2::RouteTable"
    Properties: 
      VpcId: 
        Ref: "VPC"

  PrivateSubnetRouteTableAssociation1:
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
      SubnetId:
        Ref: PrivateSubnet1
      RouteTableId:
        Ref: RouteTable

  PrivateSubnetRouteTableAssociation2:
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
      SubnetId:
        Ref: PrivateSubnet2
      RouteTableId:
        Ref: RouteTable

  # The S3 endpoint which the Jupyter notebook will use 
  # to load Braket SDK and example notebooks. This is 
  # different from the S3 endpoint in the Braket region,
  # used to store jobs and tasks' results.
  S3Endpoint:
    Type: 'AWS::EC2::VPCEndpoint'
    Properties:
      PrivateDnsEnabled: 'false'
      ServiceName: !Sub 'com.amazonaws.${AWS::Region}.s3'
      VpcEndpointType: 'Gateway'
      VpcId: !Ref VPC

  # The notebook instance that will host Braket code.
  # Since Amazon Braket is not available in the remote
  # region, you use a Sagemaker notebook instead. You will
  # customize it to run the Braket libraries later on.
  # Spinning a development environment in EC2 or Cloud9
  # are also valid alternatives to this method.

  NotebookPolicy:
    Type: AWS::IAM::Role
    Properties: 
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - braket.amazonaws.com
                - sagemaker.amazonaws.com
            Action:
              - 'sts:AssumeRole'
      Description: The role that will let the notebook interact with Amazon Braket.
      Policies: 
        - PolicyName: admin
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: Allow
                Action: '*'
                Resource: '*'

  NotebookInstance:
    Type: AWS::SageMaker::NotebookInstance
    Properties: 
      DirectInternetAccess: Disabled
      InstanceType: ml.t3.medium
      NotebookInstanceName: !Ref BraketNotebookInstance
      RoleArn: !GetAtt NotebookPolicy.Arn
      RootAccess: Enabled
      SecurityGroupIds: 
        - !GetAtt VPC.DefaultSecurityGroup
      SubnetId: !Ref PrivateSubnet1

Outputs:
  VPCId:
    Description: Local Region VPC ID
    Value: !Ref VPC

  RouteTable:
    Description: "Local Region route table ID that needs to be updated with the VPC peering connection ID."
    Value: !Ref RouteTable

  NotebookInstance:
    Description: Managed notebook from SageMaker hosting Amazon Braket code.
    Value: !Ref BraketNotebookInstance
