import os
import sys
import unittest

from torch.nn import CrossEntropyLoss
from torch.optim import SGD

from avalanche.logging import TextLogger
from avalanche.models import SimpleMLP
from avalanche.benchmarks.scenarios.online_scenario import OnlineCLScenario
from avalanche.training import OnlineNaive
from tests.unit_tests_utils import get_fast_benchmark


class StrategyTest(unittest.TestCase):
    if "FAST_TEST" in os.environ:
        fast_test = os.environ["FAST_TEST"].lower() in ["true"]
    else:
        fast_test = False
    if "USE_GPU" in os.environ:
        use_gpu = os.environ["USE_GPU"].lower() in ["true"]
    else:
        use_gpu = False

    print("Fast Test:", fast_test)
    print("Test on GPU:", use_gpu)

    if use_gpu:
        device = "cuda"
    else:
        device = "cpu"

    def init_sit(self):
        model = self.get_model(fast_test=True)
        optimizer = SGD(model.parameters(), lr=1e-3)
        criterion = CrossEntropyLoss()
        benchmark = self.load_benchmark(use_task_labels=False)
        benchmark_streams = benchmark.streams.values()
        ocl_benchmark = OnlineCLScenario(benchmark_streams)
        return model, optimizer, criterion, ocl_benchmark

    def test_naive(self):
        # SIT scenario
        model, optimizer, criterion, my_nc_benchmark = self.init_sit()
        strategy = OnlineNaive(
            model,
            optimizer,
            criterion,
            train_mb_size=1,
            device=self.device,
            eval_mb_size=50,
        )
        self.run_strategy(my_nc_benchmark, strategy)

        # MT scenario
        strategy = OnlineNaive(
            model,
            optimizer,
            criterion,
            train_mb_size=1,
            device=self.device,
            eval_mb_size=50
        )
        benchmark = self.load_benchmark(use_task_labels=True)
        benchmark_streams = benchmark.streams.values()
        ocl_benchmark = OnlineCLScenario(benchmark_streams)
        self.run_strategy(ocl_benchmark, strategy)

    def load_benchmark(self, use_task_labels=False):
        """
        Returns a NC benchmark from a fake dataset of 10 classes, 5 experiences,
        2 classes per experience.

        :param fast_test: if True loads fake data, MNIST otherwise.
        """
        return get_fast_benchmark(use_task_labels=use_task_labels)

    def get_model(self, fast_test=False):
        if fast_test:
            model = SimpleMLP(input_size=6, hidden_size=10)
            # model.classifier = IncrementalClassifier(
            #     model.classifier.in_features)
            return model
        else:
            model = SimpleMLP()
            # model.classifier = IncrementalClassifier(
            #     model.classifier.in_features)
            return model

    def run_strategy(self, benchmark, cl_strategy):
        print("Starting experiment...")
        cl_strategy.evaluator.loggers = [TextLogger(sys.stdout)]
        results = []
        for train_batch_info in benchmark.train_stream:
            print("Start of experience ", train_batch_info.current_experience)

            cl_strategy.train(train_batch_info, num_workers=0)
            print("Training completed")

            print("Computing accuracy on the current test set")
            results.append(cl_strategy.eval(benchmark.original_test_stream[:]))


if __name__ == "__main__":
    unittest.main()

@@ -1236,22 +1236,149 @@
     
-# SIT scenario
+benchmark = self.load_benchmark(use_task_labels=False)%0A        benchmark_streams = benchmark.streams.values()%0A%0A        # With task boundaries
 %0A   
@@ -1630,17 +1630,16 @@
 _size=50
-,
 %0A       
@@ -1631,32 +1631,162 @@
 size=50%0A        
+)%0A        ocl_benchmark = OnlineCLScenario(benchmark_streams,%0A                                         access_task_boundaries=True
 )%0A        self.r
@@ -1800,14 +1800,23 @@
 tegy
-(my_nc
+_boundaries(ocl
 _ben
@@ -1848,19 +1848,102 @@
   # 
-MT scenario
+Without task boundaries%0A        model, optimizer, criterion, my_nc_benchmark = self.init_sit()
 %0A   
@@ -2135,32 +2135,36 @@
       )%0A        
+ocl_
 benchmark = self
@@ -2163,163 +2163,113 @@
 k = 
-self.load_benchmark(use_task_labels=True)%0A        benchmark_streams = benchmark.streams.values()%0A        ocl_benchmark = OnlineCLScenario(benchmark_streams
+OnlineCLScenario(benchmark_streams,%0A                                         access_task_boundaries=False
 )%0A  
@@ -2291,16 +2291,30 @@
 strategy
+_no_boundaries
 (ocl_ben
@@ -3110,16 +3110,27 @@
 strategy
+_boundaries
 (self, b
@@ -3187,16 +3187,35 @@
 periment
+ (with boundaries) 
 ...%22)%0A  
@@ -3683,16 +3683,470 @@
 m%5B:%5D))%0A%0A
+    def run_strategy_no_boundaries(self, benchmark, cl_strategy):%0A        print(%22Starting experiment (without boundaries) ...%22)%0A        cl_strategy.evaluator.loggers = %5BTextLogger(sys.stdout)%5D%0A        results = %5B%5D%0A%0A        cl_strategy.train(benchmark.train_stream, num_workers=0)%0A        print(%22Training completed%22)%0A%0A        print(%22Computing accuracy on the current test set%22)%0A        results.append(cl_strategy.eval(benchmark.original_test_stream%5B:%5D))%0A%0A
 %0Aif __na

Fix `test_online_strategies`