You just make the player character fight different version of the rival character based on what Pokemon the Player character took. Take a game variable, set it to 1 if they take the grass, 2 if they take the water, and 3 if they take the fire. Then when they battle have it check that and set up a trainer battle against the appropriate version of the rival.