How to use `for_each` on a to be created resource in Terraform

I have 3 private and 3 public subnets, and I want to create either 1 or 2 NAT Gateways depending on the environment.

So, for dev and staging there should be 1 NAT Gateway, and for production two.

Since I can't use for_each with subnets as it is one more than I need, I decided to create the Elastic IPs with count.

resource "aws_eip" "elastic_ip" {  count = var.environment == "stg" ? 1 : 2  vpc = true  tags = merge(var.tags, {    Name        = "eip-${var.name_suffix}-${count.index}"    Description = "Terraform Managed Elastic IP"    Project     = var.project    Environment = var.environment  })}

Now I want to create the NAT Gateways depending on the number of Elastic IPs created, like this:

resource "aws_nat_gateway" "nat_gw" {  for_each      = aws_eip.elastic_ip  allocation_id = each.value.id  subnet_id     = CHALLENGE 2  tags = merge(var.tags, {    Name        = "nat-gw-${var.name_suffix}"    Description = "Terraform Managed NAT Gateway"    Project     = var.project    Environment = var.environment  })}

but here I have two challenges.

1. it complaints about aws_eip.elastic_ip being a tuple. I have tried using toset(), but it didn't work
2. I need to dynamically pull the IDs of two out of three subnets.

Is this actually possible without extra locals or variables?

UPDATE

Here the code for subnets. I have two of this; one for public subnets and this one for private.

resource "aws_subnet" "private_subnet" {  for_each = var.private_subnet  availability_zone = each.value["az"]  cidr_block        = each.value["cidr"]  vpc_id            = aws_vpc.vpc.id  tags = merge(var.tags, {    Name        = "private-subnet-${var.name_suffix}"    Description = "Terraform Managed Subnet"    Project     = var.project    Environment = var.environment    AZ          = each.value["az"]  })}

And I'm passing a variable with the region and cidr blocks:

private_subnet = {  subnet_a = {    az = "eu-west-1a"    cidr = "10.10.0.0/24"  }  subnet_b = {    az = "eu-west-1b"    cidr = "10.10.1.0/24"  }  subnet_c = {    az = "eu-west-1c"    cidr = "10.10.2.0/24"  }}