#include <QFile>
#include <QDebug>
#include <QApplication>
#include <QDir>
#include <QTime>
#include <QCoreApplication>
#include <QLayout>
#include <QMessageBox>
#include <QFileDialog>
#include <QMap>

#include "CSVReader.h"
#include "tinyxml2.h"
#include "annotationxml.h"
#include "object.h"
#include "newsequenz.h"

// Lädt eine Bildsequenz
//  directory: Der Pfad zum Ordner, der die Bildsequenz enthält
//  status: Im Verlauf des Ladens wird der Statustext in diesem Label
//          gesetzt
AnnotationLoader::AnnotationLoader( QString file, QLabel *status )
{
    sequenz = 0;
    error = "Pfad nicht gefunden.";
    QDir dir( file + "/Annotations" );
    QFileInfoList fileInfos = dir.entryInfoList();
    int fileCount = fileInfos.count();
    QList< Kamera* > cams;
    QList< Object > objects;
    QList< QString > classes;
    QProgressDialog progress( "Annotation wird geladen...", "", 0, fileCount );
    progress.setWindowModality(Qt::WindowModal);
    progress.setCancelButton( 0 );
    for( int i = 2; i < fileCount; i++ )
    { // Schleife durch alle xml Dateien mit Annotationen
        status->setText( "frame " + QString::number( i - 1 ) + "/" + QString::number( fileCount - 2 ) + ", found cameras: " + QString::number( cams.count() ) + ", found packets: " + QString::number( objects.count() ) + " . . ." );
        status->setGeometry( status->x(), status->y(), status->fontMetrics().width( status->text() ) + 10, status->height() );
        status->repaint(); // Aktualisiere Status Text
        tinyxml2::XMLDocument doc;
        doc.LoadFile(fileInfos.at( i ).absoluteFilePath().toStdString().c_str() );
        tinyxml2::XMLElement *camera = doc.FirstChildElement()->FirstChildElement("camera_id");
        QString camName = camera->GetText();
        Kamera *c = 0;
        for( auto cam = cams.begin(); cam != cams.end(); cam++ )
        { // Suche in den Kameras, ob die Kamera des aktuellen Bildes bereits existiert
            if( (*cam)->getName() == camName )
                c = *cam;
        }
        if( !c )
        { // Falls die Kamera noch nicht existiert wird diese erstellt
            c = new Kamera( camName, file + "/SourceMasks/" + QString::number( cams.size() ) + ".jpg", cams.size() );
            cams.append( c );
        }
        Frame *f = new Frame( file + "/SourceImages/" + doc.FirstChildElement()->FirstChildElement("filename")->GetText(),
                              doc.FirstChildElement()->FirstChildElement("timestamp")->GetText(), c->getChildCount(), c, doc.FirstChildElement()->FirstChildElement( "noobject" ) != 0 );
        c->addFrame( f );
        tinyxml2::XMLElement *object = doc.FirstChildElement()->FirstChildElement("object");
        if( object )
        {
            do
            { // Schleife durch alle annotierten Objekte des Bildes
                int truncated = 0;
                QString className = object->FirstChildElement("name")->GetText();
                int classId = classes.indexOf( className );
                if( classId < 0 )
                {
                    classId = classes.size();
                    classes.append( className );
                }
                object->FirstChildElement("truncated")->QueryIntText(&truncated);
                QList< QPolygon > polygons;
                tinyxml2::XMLElement *polygon = object->FirstChildElement( "polygon" );
                if( polygon )
                {
                    do
                    { // Schleife durch alle Polygone des Objektes
                        QPolygon pol;
                        tinyxml2::XMLElement *point = polygon->FirstChildElement( "point" );
                        if( point )
                        {
                            do
                            { // Schleife durch alle Punkte des Polygons
                                int x, y;
                                point->FirstChildElement( "x" )->QueryIntText( &x );
                                point->FirstChildElement( "y" )->QueryIntText( &y );
                                pol.push_back( QPoint( x, y ) );
                            } while( (point = point->NextSiblingElement() ) != NULL );
                            polygons.append( pol );
                        }
                    } while( (polygon = polygon->NextSiblingElement() ) != NULL );
                }
                QString pName = object->FirstChildElement("id")->GetText();
                if( objects.indexOf( Object( pName, classId ) ) < 0 )
                    objects.append( Object( pName, classId ) );
                f->addObject( pName, truncated != 0, polygons );
            } while( ( object = object->NextSiblingElement() ) != NULL );
        }
        f->wasChangedSinceLastSave();
        progress.setValue( progress.value() + 1 );
    }
    sequenz = new Sequenz( file, cams, objects );
    if( !cams.size() )
    {
        sequenz->refRelease();
        sequenz = 0;
    }
    foreach( QString className, classes )
    {
        sequenz->addClass( className );
    }
}

AnnotationLoader::~AnnotationLoader()
{
    if( sequenz )
        sequenz->refRelease();
}

// Gibt die geladene Sequenz zurück
//  0 falls beim Laden ein Fehler aufgetreten ist
Sequenz *AnnotationLoader::getSequenz()
{
    if( sequenz )
        sequenz->refNew();
    return sequenz;
}

// Gibt eine Fehlermeldung zurück, falls beim Laden ein Fehler aufgetreten
// ist
QString AnnotationLoader::getErrorMessage()
{
    return error;
}

// Erstellt eine neue Sequenz und Lädt diese
//  directory: Der Pfad zu dem Ordner, welche die Bilder der neuen Sequenz
// enthält
//  status: Im Verlauf des Erstellens wird der Statustext in diesem Label
// gesetzt
AnnotationCreator::AnnotationCreator( QString directory, QLabel *status )
{
    QDir dir( directory );
    QFileInfoList fileInfos = dir.entryInfoList();
    int fileCount = fileInfos.count();
    QString pathToCSV = "";
    for( int i = 2; i < fileCount; i++ )
    { // Suche nach der CSV Datei im angegebenen Ordner
        if( fileInfos.at( i ).suffix().toLower() == "csv" )
        {
            CSVReader reader( fileInfos.at( i ).absoluteFilePath().toStdString() );
            bool ok = false;
            while( reader.hasNext() )
            { // Suche nach dem Info Bild in der csv Datei
                std::vector< std::string > row = reader.getNextRow();
                if( row.size() < 3 )
                { // Prüfe, ob die csv Datei das passende Vormat hat
                    ok = false;
                    break;
                }
                if( row.at( 1 ) == "Info Bild" )
                    ok = true;
            }
            if( ok )
                pathToCSV = fileInfos.at( i ).absoluteFilePath();
        }
    }
    QString targetDir;
    if( pathToCSV == "" )
    { // Es wurde keine csv Datei gefunden
        QMessageBox::StandardButton reply = QMessageBox::question(0, "Achtung", "In dem ausgewählten Ordner wurde keine CSV Datei mit Informationen über die neue Sequenz gefunden. Soll der Ordner manuell nach Bildern durchsucht werden? Die Bilder würden der Sequenz in Alphabetischer Reihenfolge hinzugefügt werden.",
                                                                  QMessageBox::Yes|QMessageBox::No);
        if (reply == QMessageBox::Yes) {
            do {
                targetDir = QFileDialog::getExistingDirectory(0,
                                                              "Wo soll die neue Sequenz gespeichert werden?",
                                                              "/studi-tmp/kstrohm",
                                                              QFileDialog::ShowDirsOnly |
                                                              QFileDialog::DontResolveSymlinks);
                if( !QDir( targetDir ).count() )
                    QMessageBox::critical(0,"Fehler","Bitte wählen sie einen leeren Ordner aus.");
            } while( !QDir( targetDir ).count() );
            status->setText("Suche Nach Bildern ...");
            status->setGeometry( status->x(), status->y(), status->fontMetrics().width( status->text() ) + 10, status->height() );
            status->repaint(); // Aktualisiere Status Text
            numFrames = 0;
            for( int i = 2; i < fileCount; i++ )
            { // Zähle alle Bilder im angegebenen Ordner
                numFrames += getImageCount( fileInfos.at( i ) );
            }
            if( !numFrames )
            {
                QMessageBox::critical(0,"Fehler","Es wurden keine Bilder gefunden. Unterstützte Bildformate sind: jpg, png und bmp.");
            }
            else
            {
                NewSequenz nseq( numFrames );
                nseq.exec(); // Frage den Nutzer, welche Bilder er in die Neue Bildsequenz übernehmen möchte
                int offset = nseq.getOffset(); // Startindex des ersten Bildes
                int limit = nseq.getLimit(); // Anzahl der zu übernehmenden Bilder
                numFrames = limit;
                if( limit > 0 && offset >= 0 )
                { // Falls das erstellen nicht abgebrochen wurde
                    status->setText("Erstelle Annotation ...");
                    status->setGeometry( status->x(), status->y(), status->fontMetrics().width( status->text() ) + 10, status->height() );
                    status->repaint();
                    int index = 0;
                    // Erzteuge Ordnerstruktur
                    QDir().mkpath( targetDir + "/Annotations" );
                    QDir().mkpath( targetDir + "/SourceImages" );
                    QDir().mkpath( targetDir + "/SourceMasks" );
                    QDir().mkpath( targetDir + "/JPEGImages" );
                    int nLen = 0;
                    int count = limit;
                    while( count > 0 )
                    { // Ermittle die maximale Länge des Bildnamens (Die bilder werden 00..01-99..99 nummerriert)
                        count /= 10;
                        nLen++;
                    }
                    QProgressDialog progress( "Annotation wird erstellt...", "", 0, limit );
                    progress.setWindowModality(Qt::WindowModal);
                    progress.setCancelButton( 0 );
                    int kamId = 0;
                    for( int i = 2; i < fileCount; i++ )
                    { // Erzeuge die neue Bildsequenz
                        createAnnotationFor( fileInfos.at( i ), index, "", targetDir, nLen, status, offset, limit, progress, kamId );
                    }
                    AnnotationLoader loader( targetDir, status );
                    sequenz = loader.getSequenz(); // lade die Bildsequenz
                }
            }
        }
    }
    else
    { // Falls eine CSV Datei existiert
        do {
            targetDir = QFileDialog::getExistingDirectory(0,
                                                          "Wo soll die neue Sequenz gespeichert werden?",
                                                          "/studi-tmp/kstrohm",
                                                          QFileDialog::ShowDirsOnly |
                                                          QFileDialog::DontResolveSymlinks);
            if( !QDir( targetDir ).count() )
                QMessageBox::critical(0,"Fehler","Bitte wählen sie einen leeren Ordner aus.");
        } while( !QDir( targetDir ).count() );
        // Erzeuge die Ordnerstruktur
        QDir().mkpath( targetDir + "/Annotations" );
        QDir().mkpath( targetDir + "/SourceImages" );
        QDir().mkpath( targetDir + "/SourceMasks" );
        QDir().mkpath( targetDir + "/JPEGImages" );
        CSVReader reader( pathToCSV.toStdString() );
        status->setText("Suche Nach Bildern ...");
        int numFrames = 0;
        while( reader.hasNext() )
        { // Zahle die Anzahl der Bilder (ohne Info Bild)
            std::vector< std::string > row = reader.getNextRow();
            if(row.size() < 3 || row.at( 1 ) == "Info Bild")
                continue;
            numFrames++;
        }
        NewSequenz nseq( numFrames );
        nseq.exec(); // Frage den Nutzer welche Bilder in die neue Sequenz übernommen werden sollen
        int offset = nseq.getOffset();
        int limit = nseq.getLimit();
        int max = limit;
        if( limit > 0 && offset >= 0 )
        { // falls die Erstellung nicht abgebrochen wurde
            reader = CSVReader( pathToCSV.toStdString() );
            int nLen = 0;
            int count = limit;
            while( count > 0 )
            { // Ermittle die Maximale Länge des Bildnamens
                count /= 10;
                nLen++;
            }
            int index = 0;
            QProgressDialog progress( "Annotation wird erstellt...", "", 0, limit );
            progress.setWindowModality(Qt::WindowModal);
            progress.setCancelButton( 0 );
            int maskCount = -1;
            QMap< QString, int > cams;
            while( reader.hasNext() )
            { // Schleife durch jede Zeile der CSV Datei
                std::vector< std::string > row = reader.getNextRow();
                if(row.size() < 3 || row.at( 1 ) == "Info Bild")
                    continue;
                if( offset > 0 )
                {
                    offset--;
                    continue;
                }
                if( !limit )
                    break;
                else
                    limit--;
                QString name = QString::number( index++ );
                while( name.length() < nLen )
                    name = "0" + name;
                QString camName;
                if( !cams.contains( QString( row.at( 1 ).c_str() ) ) )
                {
                    cams.insert( QString( row.at( 1 ).c_str() ), ++maskCount );
                    camName = QString::number( maskCount );
                }
                else
                {
                    camName = QString::number( cams.take( QString( row.at( 1 ).c_str() ) ) );
                }
                QString sourceImagePath = targetDir + "/SourceImages/" + name + "." + row.at( 0 ).substr(row.at( 0 ).size() - 3).c_str();
                QString jpegImagePath = targetDir + "/JPEGImages/" + name + "." + row.at( 0 ).substr(row.at( 0 ).size() - 3).c_str();
                QString sourceMaskPath = targetDir + "/SourceMasks/" + camName + ".jpg";
                QString imagePath = dir.absolutePath() + "/" + row.at( 0 ).c_str();
                QFile( imagePath ).copy( sourceImagePath );
                QFile( imagePath ).copy( jpegImagePath );
                if( !QFileInfo(sourceMaskPath).exists() )
                { // Erzeuge die Maske der Kamera falls sie noch nicht existiert
                    QImage mask( imagePath );
                    mask.fill( 0xFFFFFFFF );
                    mask.save( sourceMaskPath );
                }
                Frame currentFrame( sourceImagePath, QString( row.at( 2 ).c_str() ), 0, 0, 1 );
                tinyxml2::XMLDocument doc; // Erstelle xml Datei
                doc.LinkEndChild( doc.NewDeclaration( "xml version=\"1.0\" " ) );
                tinyxml2::XMLElement *annotation = doc.NewElement( "annotation");
                annotation->LinkEndChild( doc.NewElement( "folder" ) )->LinkEndChild( doc.NewText( "VOC2007" ) );
                annotation->LinkEndChild( doc.NewElement( "filename" ) )->LinkEndChild( doc.NewText( currentFrame.getName().toStdString().c_str() ) );
                tinyxml2::XMLElement *source = doc.NewElement( "source" );
                source->LinkEndChild( doc.NewElement( "database" ) )->LinkEndChild( doc.NewText( "The VOC2007 Database" ) );
                source->LinkEndChild( doc.NewElement( "annotation" ) )->LinkEndChild( doc.NewText( "PASCAL VOC 2007" ) );
                source->LinkEndChild( doc.NewElement( "image" ) )->LinkEndChild( doc.NewText( currentFrame.getName().toStdString().c_str() ) );
                annotation->LinkEndChild( source );
                annotation->LinkEndChild( doc.NewElement( "owner" ) )->LinkEndChild( doc.NewElement( "name" ) )->LinkEndChild( doc.NewText( "Kolja Strohm" ) );
                annotation->LinkEndChild( doc.NewElement( "camera_id" ) )->LinkEndChild( doc.NewText( row.at( 1 ).c_str() ) );
                tinyxml2::XMLElement *size = doc.NewElement( "size" );
                QImage img = currentFrame.getImage();
                size->LinkEndChild( doc.NewElement( "width" ) )->LinkEndChild( doc.NewText( std::to_string( img.width() ).c_str() ) );
                size->LinkEndChild( doc.NewElement( "height" ) )->LinkEndChild( doc.NewText( std::to_string( img.height() ).c_str() ) );
                size->LinkEndChild( doc.NewElement( "depth" ) )->LinkEndChild( doc.NewText( std::to_string( img.depth() ).c_str() ) );
                annotation->LinkEndChild( size );
                annotation->LinkEndChild( doc.NewElement( "segmented" ) )->LinkEndChild( doc.NewText( "0" ) );
                annotation->LinkEndChild( doc.NewElement( "timestamp" ) )->LinkEndChild( doc.NewText( currentFrame.getTimestamp().toStdString().c_str() ) );
                annotation->LinkEndChild( doc.NewElement( "noobject" ) )->LinkEndChild( doc.NewText( "need Annotation" ) );
                doc.LinkEndChild( annotation );
                name = currentFrame.getName();
                doc.SaveFile( (targetDir + "/Annotations/" + name.mid( 0, name.lastIndexOf( '.' ) ) + ".xml").toStdString().c_str() );

                status->setText( "create Annotation from csv ... " + QString::number((int)((max - limit) / (float)max * 100)) + "% complete" );
                status->setGeometry( status->x(), status->y(), status->fontMetrics().width( status->text() ) + 10, status->height() );
                status->repaint();
                progress.setValue( progress.value() + 1 );
            }
            AnnotationLoader loader( targetDir, status );
            sequenz = loader.getSequenz();
        }
    }
}

// Zählt alle Bilder in einem Ordner
//  f: die Informationen zu dem Ordner
int AnnotationCreator::getImageCount( QFileInfo f )
{
    if( f.isDir() )
    {
        int count = 0;
        QFileInfoList fileInfos = QDir( f.absoluteFilePath() ).entryInfoList();
        for( QFileInfo f : fileInfos )
        {
            if( f.fileName() == "." || f.fileName() == ".." )
                continue;
            count += getImageCount( f );
        }
        return count;
    }
    if( f.suffix().toLower() == "jpg" || f.suffix().toLower() == "png" || f.suffix().toLower() == "bmp" )
        return 1;
    return 0;
}

// Erstellt Annotationen zu allen Bildern in einem Ordner
//  file: Die Information zur Datei, zu der Annotationen erstellt werden sollen
//  index: Der Index des Bildes
//  kamera: Der Name der aktuellen Kamera
//  targetDir: Der Pfad zum Ziel Ordner
//  nLen: Die maximale Länge des Namens
//  status: Ein Zeiger auf den Status Text
//  offset: Der Index des sersten Bildes der Sequenz
//  limit: Die Anzahl der Bilder in der Sequenz
//  progress: Ein Dialog mit einer Fortschrittsleiste
//  kamId: Die Id der momentanen Kamera
void AnnotationCreator::createAnnotationFor( QFileInfo file, int &index, QString kamera, QString targetDir, int nLen, QLabel *status, int &offset, int &limit, QProgressDialog &progress, int &kamId )
{
    if( limit == 0 )
        return;
    if( file.isDir() )
    { // Falls die Datei ein Ordner ist, erfolgt ein recursiver Aufruf für alle Dateien darin
        QString kamera = file.fileName();
        QDir kamDir( file.absoluteFilePath() );
        QFileInfoList kfInfos = kamDir.entryInfoList();
        for( QFileInfo file : kfInfos )
        {
            QString newKam = kamera;
            if( newKam != "" )
                newKam += "/";
            newKam += file.fileName();
            createAnnotationFor( file, index, newKam, targetDir, nLen, status, offset, limit, progress, kamId );
        }
        kamId++;
    }
    else
    { // Falls die Datei kein Ordner ist
        if( offset > 0 )
        {
            offset--;
            return;
        }
        if( limit > 0 )
            limit--;
        if( file.suffix().toLower() == "jpg" || file.suffix().toLower() == "png" || file.suffix().toLower() == "bmp" )
        { // Falls die Datei ein Bild ist, wird eine Annotation für sie erstellt
            QString name = QString::number( index );
            while( name.length() < nLen )
                name = "0" + name;
            if( kamera == "" )
                kamera = "Ohne Kamera";
            QString sourceImagePath = targetDir + "/SourceImages/" + name + +"." + file.suffix().toLower();
            QString jpegImagePath = targetDir + "/JPEGImages/" + name + +"." + file.suffix().toLower();
            QString sourceMaskPath = targetDir + "/SourceMasks/" + QString::number( kamId ) + +".jpg";
            QFile( file.absoluteFilePath() ).copy( sourceImagePath );
            QFile( file.absoluteFilePath() ).copy( jpegImagePath );
            if( !QFileInfo(sourceMaskPath).exists() )
            { // Falls die Maske der Kamera noch nicht existiert wird sie erstellt
                QImage mask( file.absoluteFilePath() );
                mask.fill( 0xFFFFFFFF );
                mask.save( sourceMaskPath );
            }
            Frame currentFrame( sourceImagePath, file.created().toString(Qt::SystemLocaleLongDate), 0, 0, 1 );
            tinyxml2::XMLDocument doc;
            doc.LinkEndChild( doc.NewDeclaration( "xml version=\"1.0\" " ) );
            tinyxml2::XMLElement *annotation = doc.NewElement( "annotation");
            annotation->LinkEndChild( doc.NewElement( "folder" ) )->LinkEndChild( doc.NewText( "VOC2007" ) );
            annotation->LinkEndChild( doc.NewElement( "filename" ) )->LinkEndChild( doc.NewText( currentFrame.getName().toStdString().c_str() ) );
            tinyxml2::XMLElement *source = doc.NewElement( "source" );
            source->LinkEndChild( doc.NewElement( "database" ) )->LinkEndChild( doc.NewText( "The VOC2007 Database" ) );
            source->LinkEndChild( doc.NewElement( "annotation" ) )->LinkEndChild( doc.NewText( "PASCAL VOC 2007" ) );
            source->LinkEndChild( doc.NewElement( "image" ) )->LinkEndChild( doc.NewText( currentFrame.getName().toStdString().c_str() ) );
            annotation->LinkEndChild( source );
            annotation->LinkEndChild( doc.NewElement( "owner" ) )->LinkEndChild( doc.NewElement( "name" ) )->LinkEndChild( doc.NewText( "Kolja Strohm" ) );
            annotation->LinkEndChild( doc.NewElement( "camera_id" ) )->LinkEndChild( doc.NewText( kamera.toStdString().c_str() ) );
            tinyxml2::XMLElement *size = doc.NewElement( "size" );
            QImage img = currentFrame.getImage();
            size->LinkEndChild( doc.NewElement( "width" ) )->LinkEndChild( doc.NewText( std::to_string( img.width() ).c_str() ) );
            size->LinkEndChild( doc.NewElement( "height" ) )->LinkEndChild( doc.NewText( std::to_string( img.height() ).c_str() ) );
            size->LinkEndChild( doc.NewElement( "depth" ) )->LinkEndChild( doc.NewText( std::to_string( img.depth() ).c_str() ) );
            annotation->LinkEndChild( size );
            annotation->LinkEndChild( doc.NewElement( "segmented" ) )->LinkEndChild( doc.NewText( "0" ) );
            annotation->LinkEndChild( doc.NewElement( "timestamp" ) )->LinkEndChild( doc.NewText( currentFrame.getTimestamp().toStdString().c_str() ) );
            annotation->LinkEndChild( doc.NewElement( "noobject" ) )->LinkEndChild( doc.NewText( "need Annotation" ) );
            doc.LinkEndChild( annotation );
            doc.SaveFile( (targetDir + "/Annotations/" + name + ".xml").toStdString().c_str() );
            status->setText( "create Annotation without csv ... " + QString::number((int)((numFrames - limit) / (float)numFrames * 100)) + "% complete" );
            status->setGeometry( status->x(), status->y(), status->fontMetrics().width( status->text() ) + 10, status->height() );
            status->repaint();
            index++;
            progress.setValue( progress.value() + 1 );
        }
    }
}

AnnotationCreator::~AnnotationCreator()
{
    if( sequenz )
        sequenz->refRelease();
}

// Gibt die geladene neu erstellte Sequenz zurück
//  0 falls beim Erstellen oder beim Laden ein Fehler aufgetreten ist
Sequenz *AnnotationCreator::getSequenz()
{
    if( sequenz )
        sequenz->refNew();
    return sequenz;
}