#!/usr/bin/env python
# encoding: utf-8

# The MIT License (MIT)

# Copyright (c) 2016-2018 CNRS

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

# AUTHORS
# Hervé BREDIN - http://herve.niderb.fr


import os
import os.path
import torch
from pyannote.audio.util import mkdir_p


class Checkpoint(object):
    """

    Parameters
    ----------
    log_dir : str
    restart : boolean, optional
        Indicates that this training is a restart, not a cold start (default).
    """

    WEIGHTS_DIR = '{log_dir}/weights'
    WEIGHTS_PT = '{log_dir}/weights/{epoch:04d}.pt'
    OPTIMIZER_PT = '{log_dir}/weights/{epoch:04d}.optimizer.pt'

    def __init__(self, log_dir, restart=False):
        super(Checkpoint, self).__init__()

        # make sure path is absolute
        self.log_dir = os.path.realpath(log_dir)

        # create log_dir directory
        mkdir_p(self.log_dir)

        # this will fail if the directory already exists
        # and this is OK  because 'weights' directory
        # usually contains the output of very long computations
        # and you do not want to erase them by mistake :/
        self.restart = restart
        if not self.restart:
            weights_dir = self.WEIGHTS_DIR.format(log_dir=self.log_dir)
            os.makedirs(weights_dir)

    def weights_dir(self):
        """Return path to weights directory"""
        return self.WEIGHTS_DIR.format(log_dir=self.log_dir)

    def weights_pt(self, epoch):
        """Return path to weights file"""
        return self.WEIGHTS_PT.format(log_dir=self.log_dir, epoch=epoch)

    def optimizer_pt(self, epoch):
        """Return path to optimizer file"""
        return self.OPTIMIZER_PT.format(log_dir=self.log_dir, epoch=epoch)

    def on_epoch_end(self, epoch, model, optimizer, extra=None):

        torch.save(model.state_dict(), self.weights_pt(epoch))
        torch.save(optimizer.state_dict(), self.optimizer_pt(epoch))

        if extra is not None:
            for extra_pt, extra_module in extra.items():
                path = extra_pt.format(log_dir=self.log_dir, epoch=epoch)
                torch.save(extra_module.state_dict(), path)

@@ -1306,24 +1306,41 @@
 ct):%0A    %22%22%22
+Model checkpoints
 %0A%0A    Parame
@@ -1373,19 +1373,56 @@
 g_dir : 
+%60
 str
+%60%0A        Path to logging directory.
 %0A    res
@@ -1432,15 +1432,14 @@
 t : 
+%60
 bool
-ean
+%60
 , op
@@ -1925,32 +1925,189 @@
 (self.log_dir)%0A%0A
+        self.restart = restart%0A        if not self.restart:%0A            weights_dir = self.WEIGHTS_DIR.format(log_dir=self.log_dir)%0A            try:%0A        
         # this w
@@ -2151,24 +2151,32 @@
 sts%0A        
+   
+     
 # and this i
@@ -2201,32 +2201,40 @@
 ghts' directory%0A
+        
         # usuall
@@ -2285,24 +2285,32 @@
 ons%0A        
+     
+   
 # and you do
@@ -2359,60 +2359,185 @@
     
-self.restart = restart%0A        if not self.restart:%0A
+        os.makedirs(weights_dir)%0A            except FileExistsError as e:%0A                msg = (%0A                    f'You are about to overwrite pretrained models in '%0A       
     
@@ -2536,32 +2536,37 @@
                 
+ f'%22%7B
 weights_dir = se
@@ -2564,92 +2564,218 @@
 _dir
- = self.WEIGHTS_DIR.format(log_dir=self.log_dir)%0A            os.makedirs(weights_dir
+%7D%22 directory. If you want to train a new '%0A                    f'model from scratch, first (backup and) remove the '%0A                    f'directory.'%0A                )%0A                raise FileExistsError(msg
 )%0A%0A 

make FileExistsError message more explicit