---
AWSTemplateFormatVersion: '2010-09-09'

Description: Creates an S3 Bucket, IAM Policies, and SageMaker Notebook to work with Amazon Forecast.

Parameters:

  BucketName:
    Type: String
    Default: ""
    Description: The (globally unique) name of the S3 Bucket to create, or blank to set a name automatically.
    AllowedPattern: "^[0-9a-zA-Z]+([0-9a-zA-Z-]*[0-9a-zA-Z])*$"
    ConstraintDescription": Can include numbers, letters, and hyphens (-). Must not start or end with a hyphen.

  NotebookName:
    Type: String
    Default: ForecastDemoLab
    Description: Enter the name of the SageMaker notebook instance. Deafault is ForecastDemoLab.

  VolumeSize:
    Type: Number
    Default: 10
    MinValue: 5
    MaxValue: 16384
    ConstraintDescription: Must be an integer between 5 (GB) and 16384 (16 TB).
    Description: Enter the size of the EBS volume in GB.
  
  UseVPC:
    Type: String
    AllowedValues:
      - "true"
      - "false"
      
  SecurityGroup:
    # Use String rather than AWS::EC2::SecurityGroup::Id to allow for empty values
    Type: String
    Description: The security group to provide the notebook access within the VPC
  
  PrivateSubnet:
    # Use String rather than AWS::EC2::Subnet::Id to allow for empty values
    Type: String
    Description: The ID of the subnet in a VPC to which you would like to have a connectivity from your ML compute instance

Metadata:
  AWS::CloudFormation::Interface:
    ParameterLabels:
      BucketName:
        default: S3 Bucket Name
      UseVPC:
        default: Use existing VPC?
      SecurityGroup:
        default: Security Group Id
      PrivateSubnet:
        default: Private Subnet Id
      NotebookName:
        default: Notebook Name
      VolumeSize:
        default: EBS (Hard disk) Volume Size

    ParameterGroups:
      - Label:
          default: Datastorage
        Parameters:
          - BucketName
      - Label:
          default: VPC Settings (Required if UseVPC is true)
        Parameters:
          - UseVPC
          - SecurityGroup
          - PrivateSubnet

Conditions:
  ExplicitBucketName: !Not [!Equals [!Ref BucketName, ""]]
  UseOwnVPC: !Equals [!Ref UseVPC, "true"]
  
Resources:
  # S3 First
  S3Bucket:
    Type: AWS::S3::Bucket
    Properties:
      BucketName: !If [ ExplicitBucketName, !Ref BucketName, !Ref "AWS::NoValue" ]
      BucketEncryption:
        ServerSideEncryptionConfiguration:
          -
            ServerSideEncryptionByDefault:
              SSEAlgorithm: AES256
      PublicAccessBlockConfiguration:
        BlockPublicAcls: true
        IgnorePublicAcls: true
        BlockPublicPolicy: true
        RestrictPublicBuckets: true

  # SageMaker Execution Role
  SageMakerIamRole:
    Type: "AWS::IAM::Role"
    Properties:
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          -
            Effect: Allow
            Principal:
              Service: sagemaker.amazonaws.com
            Action: sts:AssumeRole
      Path: "/"
      ManagedPolicyArns:
        - "arn:aws:iam::aws:policy/AmazonSageMakerFullAccess"
        - "arn:aws:iam::aws:policy/AmazonS3FullAccess"
        - "arn:aws:iam::aws:policy/AmazonForecastFullAccess"
        - "arn:aws:iam::aws:policy/IAMFullAccess"

  # Lifecycle configuration to automate some extra notebook setup
  NotebookConfig:
    Type: "AWS::SageMaker::NotebookInstanceLifecycleConfig"
    Properties:
      OnStart:
        - Content:
            Fn::Base64: !Sub |
              #!/bin/bash
              set -e

              # JupyterLab requires an extension to support interactive widgets used in the NBs:
              sudo -u ec2-user -i <<'EOF'
              source /home/ec2-user/anaconda3/bin/activate JupyterSystemEnv
              jupyter labextension install @jupyter-widgets/jupyterlab-manager
              source /home/ec2-user/anaconda3/bin/deactivate
              EOF

  # SageMaker notebook
  NotebookInstance:
    Type: "AWS::SageMaker::NotebookInstance"
    Properties:
      InstanceType: "ml.t2.medium"
      NotebookInstanceName: !Ref NotebookName
      LifecycleConfigName: !GetAtt NotebookConfig.NotebookInstanceLifecycleConfigName
      RoleArn: !GetAtt SageMakerIamRole.Arn
      VolumeSizeInGB: !Ref VolumeSize
      DefaultCodeRepository: https://github.com/aws-samples/amazon-forecast-samples.git
      SecurityGroupIds:
        - !If [ UseOwnVPC, !Ref SecurityGroup, !Ref "AWS::NoValue"]
      SubnetId: !If [ UseOwnVPC, !Ref PrivateSubnet, !Ref "AWS::NoValue"]

Outputs:
  S3Bucket:
    Value: !Ref S3Bucket
    Description: S3 Bucket for object storage