import json
import unittest.mock as mock

import yaml
import pytest
import tensorflow as tf

from mlagents.trainers.trainer_controller import TrainerController
from mlagents.trainers.buffer import Buffer
from mlagents.trainers.ppo.trainer import PPOTrainer
from mlagents.trainers.bc.trainer import BehavioralCloningTrainer
from mlagents.trainers.curriculum import Curriculum
from mlagents.trainers.exception import CurriculumError
from mlagents.envs.exception import UnityEnvironmentException
from tests.mock_communicator import MockCommunicator


@pytest.fixture
def dummy_start():
  return '''{ "AcademyName": "RealFakeAcademy",
              "resetParameters": {},
              "brainNames": ["RealFakeBrain"],
              "externalBrainNames": ["RealFakeBrain"],
              "logPath":"RealFakePath",
              "apiNumber":"API-5",
              "brainParameters": [{
                  "vectorObservationSize": 3,
                  "numStackedVectorObservations" : 2,
                  "vectorActionSize": 2,
                  "memorySize": 0,
                  "cameraResolutions": [],
                  "vectorActionDescriptions": ["",""],
                  "vectorActionSpaceType": 1
                  }]
            }'''.encode()


@pytest.fixture
def dummy_config():
    return yaml.load(
        '''
        default:
            trainer: ppo
            batch_size: 32
            beta: 5.0e-3
            buffer_size: 512
            epsilon: 0.2
            gamma: 0.99
            hidden_units: 128
            lambd: 0.95
            learning_rate: 3.0e-4
            max_steps: 5.0e4
            normalize: true
            num_epoch: 5
            num_layers: 2
            time_horizon: 64
            sequence_length: 64
            summary_freq: 1000
            use_recurrent: false
            memory_size: 8
            use_curiosity: false
            curiosity_strength: 0.0
            curiosity_enc_size: 1
        ''')


@pytest.fixture
def dummy_bc_config():
    return yaml.load(
        '''
        default:
            trainer: imitation
            brain_to_imitate: ExpertBrain
            batches_per_epoch: 16
            batch_size: 32
            beta: 5.0e-3
            buffer_size: 512
            epsilon: 0.2
            gamma: 0.99
            hidden_units: 128
            lambd: 0.95
            learning_rate: 3.0e-4
            max_steps: 5.0e4
            normalize: true
            num_epoch: 5
            num_layers: 2
            time_horizon: 64
            sequence_length: 64
            summary_freq: 1000
            use_recurrent: false
            memory_size: 8
            use_curiosity: false
            curiosity_strength: 0.0
            curiosity_enc_size: 1
        ''')

@pytest.fixture
def dummy_bad_config():
    return yaml.load(
        '''
        default:
            trainer: incorrect_trainer
            brain_to_imitate: ExpertBrain
            batches_per_epoch: 16
            batch_size: 32
            beta: 5.0e-3
            buffer_size: 512
            epsilon: 0.2
            gamma: 0.99
            hidden_units: 128
            lambd: 0.95
            learning_rate: 3.0e-4
            max_steps: 5.0e4
            normalize: true
            num_epoch: 5
            num_layers: 2
            time_horizon: 64
            sequence_length: 64
            summary_freq: 1000
            use_recurrent: false
            memory_size: 8
        ''')


@mock.patch('mlagents.envs.UnityEnvironment.executable_launcher')
@mock.patch('mlagents.envs.UnityEnvironment.get_communicator')
def test_initialization(mock_communicator, mock_launcher):
    mock_communicator.return_value = MockCommunicator(
        discrete_action=True, visual_inputs=1)
    tc = TrainerController(' ', ' ', 1, None, True, True, False, 1,
                           1, 1, 1, '', "tests/test_mlagents.trainers.py", False)
    assert(tc.env.brain_names[0] == 'RealFakeBrain')


@mock.patch('mlagents.envs.UnityEnvironment.executable_launcher')
@mock.patch('mlagents.envs.UnityEnvironment.get_communicator')
def test_load_config(mock_communicator, mock_launcher, dummy_config):
    open_name = 'mlagents.trainers.trainer_controller' + '.open'
    with mock.patch('yaml.load') as mock_load:
        with mock.patch(open_name, create=True) as _:
            mock_load.return_value = dummy_config
            mock_communicator.return_value = MockCommunicator(
                discrete_action=True, visual_inputs=1)
            mock_load.return_value = dummy_config
            tc = TrainerController(' ', ' ', 1, None, True, True, False, 1,
                                       1, 1, 1, '','', False)
            config = tc._load_config()
            assert(len(config) == 1)
            assert(config['default']['trainer'] == "ppo")


@mock.patch('mlagents.envs.UnityEnvironment.executable_launcher')
@mock.patch('mlagents.envs.UnityEnvironment.get_communicator')
def test_initialize_trainers(mock_communicator, mock_launcher, dummy_config,
                             dummy_bc_config, dummy_bad_config):
    open_name = 'mlagents.trainers.trainer_controller' + '.open'
    with mock.patch('yaml.load') as mock_load:
        with mock.patch(open_name, create=True) as _:
            mock_communicator.return_value = MockCommunicator(
                discrete_action=True, visual_inputs=1)
            tc = TrainerController(' ', ' ', 1, None, True, True, False, 1, 1,
                                   1, 1, '', "tests/test_mlagents.trainers.py",
                                   False)

            # Test for PPO trainer
            mock_load.return_value = dummy_config
            config = tc._load_config()
            tf.reset_default_graph()
            with tf.Session() as sess:
                tc._initialize_trainers(config, sess)
                assert(len(tc.trainers) == 1)
                assert(isinstance(tc.trainers['RealFakeBrain'], PPOTrainer))

            # Test for Behavior Cloning Trainer
            mock_load.return_value = dummy_bc_config
            config = tc._load_config()
            tf.reset_default_graph()
            with tf.Session() as sess:
                tc._initialize_trainers(config, sess)
                assert(isinstance(tc.trainers['RealFakeBrain'], BehavioralCloningTrainer))

            # Test for proper exception when trainer name is incorrect
            mock_load.return_value = dummy_bad_config
            config = tc._load_config()
            tf.reset_default_graph()
            with tf.Session() as sess:
                with pytest.raises(UnityEnvironmentException):
                    tc._initialize_trainers(config, sess)