Description: "Create Quantinuum InQuanto Notebook instance"
Parameters:
  InstanceName:
    Type: String
    Default: quantinuum-notebook
    MinLength: 3
    MaxLength: 40
    AllowedPattern: ^[a-zA-Z0-9\-]*$
    Description: Enter the notebook name - Only alphanumerals and dashes allowed
  InstanceType:
    Type: String
    Default: ml.t3.medium
    AllowedValues:
      - ml.t3.medium
      - ml.m5.xlarge
      - ml.r5.2xlarge
    Description: Enter ml.t3.medium, ml.m5.xlarge, or ml.r5.2xlarge. Default is ml.t3.medium.
Resources:
  defaultRepository:
    Type: AWS::SageMaker::CodeRepository
    Properties: 
      CodeRepositoryName: 
        !Join [ "-",[Ref: InstanceName,"repo"]]
      GitConfig:
        RepositoryUrl:  <User git repository URL>
  BasicNotebookInstanceLifecycleConfig:
    Type: "AWS::SageMaker::NotebookInstanceLifecycleConfig"
    Properties:
      NotebookInstanceLifecycleConfigName: 
        !Join [ "-",[Ref: InstanceName,"LC"]]
      OnStart:
      - Content:
          Fn::Base64: |
            #!/bin/bash
            set -e
            sudo -u ec2-user -i <<'EOF'
            unset SUDO_UID
            WORKING_DIR=/home/ec2-user/SageMaker/custom-miniconda/
            source "$WORKING_DIR/miniconda/bin/activate"
            for env in $WORKING_DIR/miniconda/envs/*; do
                BASENAME=$(basename "$env")
                source activate "$BASENAME"
                python -m ipykernel install --user --name "$BASENAME" --display-name "Custom ($BASENAME)"
            done
            # Optionally, uncomment these lines to disable SageMaker-provided Conda functionality.
            echo "c.EnvironmentKernelSpecManager.use_conda_directly = False" >> /home/ec2-user/.jupyter/jupyter_notebook_config.py
            rm /home/ec2-user/.condarc
            EOF

            echo "Restarting the Jupyter server.."
            restart jupyter-server
      OnCreate:
      - Content:
          Fn::Base64: |
            #!/bin/bash
            set -e
            sudo -u ec2-user -i <<'EOF'
            unset SUDO_UID
            # Install a separate conda installation via Miniconda
            WORKING_DIR=/home/ec2-user/SageMaker/custom-miniconda
            mkdir -p "$WORKING_DIR"
            wget https://repo.anaconda.com/miniconda/Miniconda3-4.6.14-Linux-x86_64.sh -O "$WORKING_DIR/miniconda.sh"
            bash "$WORKING_DIR/miniconda.sh" -b -u -p "$WORKING_DIR/miniconda" 
            rm -rf "$WORKING_DIR/miniconda.sh"

            # Create a custom conda environment
            source "$WORKING_DIR/miniconda/bin/activate"
            KERNEL_NAME="Quantinuum-Kernel"
            PYTHON="3.9"
            conda create --yes --name "$KERNEL_NAME" python="$PYTHON"
            conda activate "$KERNEL_NAME"
            pip install --quiet ipykernel
            # Customize these lines as necessary to install the required packages
            conda install --yes numpy
            pip install "inquanto[ibm, quantinuum, jupyter, qulacs]" inquanto-pyscf inquanto-nglview -i <Quantinuum index URL>
            EOF
  ExecutionRole: 
    Type: "AWS::IAM::Role"
    Properties: 
      RoleName: !Join [ "-", ["quantinuum-notebook-role", !Select [2, !Split [ "/", !Ref 'AWS::StackId']]]]
      AssumeRolePolicyDocument: 
        Version: "2012-10-17"
        Statement: 
        - 
          Effect: "Allow"
          Principal: 
            Service: 
              - "sagemaker.amazonaws.com"
          Action: 
          - "sts:AssumeRole"
      Path: "/"
      Policies: 
        - 
          PolicyName: !Join [ "-", ["quantinuum-notebook-policy", !Select [2, !Split [ "/", !Ref 'AWS::StackId']]]]
          PolicyDocument: 
            Version: "2012-10-17"
            Statement: 
              - Effect: "Allow"
                Action: 
                  - "cloudwatch:PutMetricData"
                  - "logs:CreateLogStream"
                  - "logs:PutLogEvents"
                  - "logs:CreateLogGroup"
                  - "logs:DescribeLogStreams"
                Resource: "*"
  sageMakerNotebook:
    Type: AWS::SageMaker::NotebookInstance
    Properties: 
      DefaultCodeRepository: !GetAtt defaultRepository.CodeRepositoryName
      DirectInternetAccess: Enabled
      InstanceType: 
        Ref: InstanceType
      NotebookInstanceName: 
        Ref: InstanceName
      RoleArn: !GetAtt ExecutionRole.Arn
      VolumeSizeInGB: 50
      LifecycleConfigName: !GetAtt BasicNotebookInstanceLifecycleConfig.NotebookInstanceLifecycleConfigName
Outputs:
  sageMakerNotebookId:
    Value: !Ref sageMakerNotebook
  sageMakerNotebookLink:
    Value: !Join [ "", ["https://",!Ref "AWS::Region",".console.aws.amazon.com/sagemaker/home?region=",!Ref "AWS::Region","#/notebook-instances/openNotebook/",!Select [1, !Split [ "/", !Ref sageMakerNotebook]],"?view=classic"]]
    
