#include "Ogre.h"

void gridScreenshots(Ogre::RenderWindow* pRenderWindow, Ogre::Camera* pCamera, const int& pGridSize, const Ogre::String& pFileName, const Ogre::String& pFileExtention, const bool& pStitchGridImages)
{
  /* Parameters:
   *  pRenderWindow:    Pointer to the render window.  This could be "mWindow" from the ExampleApplication,
   *              the window automatically created obtained when calling
   *              Ogre::Root::getSingletonPtr()->initialise(false) and retrieved by calling
   *              "Ogre::Root::getSingletonPtr()->getAutoCreatedWindow()", or the manually created
   *              window from calling "mRoot->createRenderWindow()".
   *  pCamera:      Pointer to the camera "looking at" the scene of interest
   *  pGridSize:      The magnification factor.  A 2 will create a 2x2 grid, doubling the size of the
                screenshot.  A 3 will create a 3x3 grid, tripling the size of the screenshot.
   *  pFileName:      The filename to generate, without an extention.  To generate "MyScreenshot.png" this
   *              parameter would contain the value "MyScreenshot".
   *  pFileExtention:    The extention of the screenshot file name, hence the type of graphics file to generate.
   *              To generate "MyScreenshot.png" this parameter would contain ".png".
   *  pStitchGridImages:  Determines whether the grid screenshots are (true) automatically stitched into a single
   *              image (and discarded) or whether they should (false) remain in their unstitched
   *              form.  In that case they are sequentially numbered from 0 to
   *              pGridSize * pGridSize - 1 (if pGridSize is 3 then from 0 to 8).
   *              
  */
  Ogre::String gridFilename;

  if(pGridSize <= 1)
  {
    // Simple case where the contents of the screen are taken directly
    // Also used when an invalid value is passed within pGridSize (zero or negative grid size)
    gridFilename = pFileName + pFileExtention;

    pRenderWindow->writeContentsToFile(gridFilename);
  }
  else
  {
    // Generate a grid of screenshots
    pCamera->setCustomProjectionMatrix(false); // reset projection matrix
    Ogre::Matrix4 standard = pCamera->getProjectionMatrix();
    double nearDist = pCamera->getNearClipDistance();
    double nearWidth = (pCamera->getWorldSpaceCorners()[0] - pCamera->getWorldSpaceCorners()[1]).length();
    double nearHeight = (pCamera->getWorldSpaceCorners()[1] - pCamera->getWorldSpaceCorners()[2]).length();
    Ogre::Image sourceImage;
    Ogre::uchar* stitchedImageData;

    // Process each grid
    for (int nbScreenshots = 0; nbScreenshots < pGridSize * pGridSize; nbScreenshots++) 
    { 
      // Use asymmetrical perspective projection. For more explanations check out:
      // http://www.cs.kuleuven.ac.be/cwis/research/graphics/INFOTEC/viewing-in-3d/node8.html 
      int y = nbScreenshots / pGridSize; 
      int x = nbScreenshots - y * pGridSize; 
      Ogre::Matrix4 shearing( 
        1, 0,(x - (pGridSize - 1) * 0.5) * nearWidth / nearDist, 0, 
        0, 1, -(y - (pGridSize - 1) * 0.5) * nearHeight / nearDist, 0, 
        0, 0, 1, 0, 
        0, 0, 0, 1); 
      Ogre::Matrix4 scale( 
        pGridSize, 0, 0, 0, 
        0, pGridSize, 0, 0, 
        0, 0, 1, 0, 
        0, 0, 0, 1); 
      pCamera->setCustomProjectionMatrix(true, standard * shearing * scale);
      Ogre::Root::getSingletonPtr()->renderOneFrame();
      gridFilename = pFileName + Ogre::StringConverter::toString(nbScreenshots) + pFileExtention;

      // Screenshot of the current grid
      pRenderWindow->writeContentsToFile(gridFilename);

      if(pStitchGridImages)
      {
        // Automatically stitch the grid screenshots
        sourceImage.load(gridFilename, "General"); // Assumes that the current directory is within the "General" resource group
        int sourceWidth = (int) sourceImage.getWidth();
        int sourceHeight = (int) sourceImage.getHeight();
        Ogre::ColourValue colourValue;
        int stitchedX, stitchedY, stitchedIndex;

        // Allocate memory for the stitched image when processing the screenshot of the first grid
        if(nbScreenshots == 0)
          stitchedImageData = new Ogre::uchar[(sourceImage.getWidth() * pGridSize) * (sourceImage.getHeight() * pGridSize) * 3]; // 3 colors per pixel

        // Copy each pixel within the grid screenshot to the proper position within the stitched image
        for(int rawY = 0; rawY < sourceHeight; rawY++)
        {
          for(int rawX = 0; rawX < sourceWidth; rawX++)
          {
            colourValue = sourceImage.getColourAt(rawX, rawY, 0);
            stitchedX = x * sourceWidth + rawX;
            stitchedY = y * sourceHeight + rawY;
            stitchedIndex = stitchedY * sourceWidth * pGridSize + stitchedX;
            Ogre::PixelUtil::packColour(colourValue,
                          Ogre::PF_R8G8B8,
                          (void*) &stitchedImageData[stitchedIndex * 3]);
          }
        }
        // The screenshot of the grid is no longer needed
        remove(gridFilename.c_str());
      }
    } 
    pCamera->setCustomProjectionMatrix(false); // reset projection matrix 

    if(pStitchGridImages)
    {
      // Save the stitched image to a file
      Ogre::Image targetImage;
      targetImage.loadDynamicImage(stitchedImageData,
                    sourceImage.getWidth() * pGridSize,
                    sourceImage.getHeight() * pGridSize,
                    1, // depth
                    Ogre::PF_R8G8B8,
                    false);
      targetImage.save(pFileName + pFileExtention);
      delete[] stitchedImageData;
    }
  }
}

#ifndef __ScreenshotManager_h__
#define __ScreenshotManager_h__

   /* Class encapsulates Screenshot functionality and provides a method for making multi grid screenshots.    
   *  pRenderWindow:    Pointer to the render window.  This could be "mWindow" from the ExampleApplication,
   *              the window automatically created obtained when calling
   *              Ogre::Root::getSingletonPtr()->initialise(false) and retrieved by calling
   *              "Ogre::Root::getSingletonPtr()->getAutoCreatedWindow()", or the manually created
   *              window from calling "mRoot->createRenderWindow()".
   *  gridSize:      The magnification factor.  A 2 will create a 2x2 grid, doubling the size of the
                screenshot.  A 3 will create a 3x3 grid, tripling the size of the screenshot.
   *  fileExtension:    The extension of the screenshot file name, hence the type of graphics file to generate.
   *              To generate "MyScreenshot.png" this parameter would contain ".png".
   */
    class ScreenshotManager
    {
    public:
      ScreenshotManager(Ogre::RenderWindow* pRenderWindow, int gridSize, Ogre::String fileExtension, bool overlayFlag);
      ~ScreenshotManager();

   /* Creates a screenshot with the given camera.
    * @param camera Pointer to the camera "looking at" the scene of interest
    * @param fileName the filename of the screenshot file.
   */
      void makeScreenshot(Ogre::Camera* camera, Ogre::String fileName) const;
     
   protected:
      Ogre::String     mFileExtension;
      unsigned int   mGridSize, mWindowWidth, mWindowHeight;
      bool           mDisableOverlays;
      //temp texture with current screensize
      Ogre::TexturePtr mTempTex;
      Ogre::RenderTexture* mRT;
      Ogre::HardwarePixelBufferSharedPtr mBuffer;
      //PixelBox for a large Screenshot, if grid size is > 1
      Ogre::PixelBox  mFinalPicturePB;
      //Pointer to the color data of the pixel box
      Ogre::uint8* mData;
    };

#endif  // __ScreenshotManager_h__

#include <stdafx.h>
#include "ScreenshotManager.h"

using namespace Ogre;

ScreenshotManager::ScreenshotManager(Ogre::RenderWindow* pRenderWindow, int gridSize, Ogre::String fileExtension, bool overlayFlag)
{
    //set file extension for the Screenshot files
    mFileExtension = fileExtension;
    // the gridsize
    mGridSize         = gridSize;
    // flag for overlay rendering
    mDisableOverlays  = overlayFlag;
    //get current window size
    mWindowWidth     = pRenderWindow->getWidth();
    mWindowHeight     = pRenderWindow->getHeight();
    //create temporary texture
    mTempTex    = TextureManager::getSingleton().createManual("ScreenShotTex", 
                                                               ResourceGroupManager::DEFAULT_RESOURCE_GROUP_NAME, TEX_TYPE_2D,
                                                                 mWindowWidth, mWindowHeight,0, PF_B8G8R8, TU_RENDERTARGET); 
    
    //get The current Render Target of the temp Texture
    mRT = mTempTex->getBuffer()->getRenderTarget();

    //HardwarePixelBufferSharedPtr to the Buffer of the temp Texture 
    mBuffer = mTempTex->getBuffer();

    //create PixelBox
    mData = new uint8[(mWindowWidth * mGridSize) * (mWindowHeight * mGridSize) * 3];
    mFinalPicturePB = PixelBox(mWindowWidth * mGridSize,mWindowHeight * mGridSize,1,PF_B8G8R8,mData);

}

ScreenshotManager::~ScreenshotManager()
{
   delete[] mData;
}

/* Creates a screenshot with the given camera.
* @param camera Pointer to the camera "looking at" the scene of interest
* @param fileName the filename of the screenshot file.
*/
void ScreenshotManager::makeScreenshot(Ogre::Camera* camera, Ogre::String fileName) const
{

    //Remove all viewports, so the added Viewport(camera) ist the only 
    mRT->removeAllViewports();
    mRT->addViewport(camera);
            
    //set the viewport settings
    Viewport *vp = mRT->getViewport(0);
    vp->setClearEveryFrame(true);    
    vp->setOverlaysEnabled(false);

    // remind current overlay flag
    bool enableOverlayFlag = myApplication::getSingletonPtr()->getDefaultViewport()->getOverlaysEnabled();

    // we disable overlay rendering if it is set in config file and the viewport setting is enabled
    if(mDisableOverlays && enableOverlayFlag)
        MyApplication::getSingletonPtr()->getDefaultViewport()->setOverlaysEnabled(false);

    if(mGridSize <= 1)
    {
        // Simple case where the contents of the screen are taken directly
        // Also used when an invalid value is passed within gridSize (zero or negative grid size)
        mRT->update();        //render

        //write the file on the Harddisk
        mRT->writeContentsToFile(fileName + "." + mFileExtension);
    }
    else
    {
        //define the original frustum extents variables
        Real originalFrustumLeft, originalFrustumRight, originalFrustumTop, originalFrustumBottom;
        // set the original Frustum extents
        camera->getFrustumExtents(originalFrustumLeft, originalFrustumRight, originalFrustumTop, originalFrustumBottom);
        
        // compute the Stepsize for the drid
        Real frustumGridStepHorizontal    = (originalFrustumRight * 2) / mGridSize;
        Real frustumGridStepVertical    = (originalFrustumTop * 2) / mGridSize;

        // process each grid
        Real frustumLeft, frustumRight, frustumTop, frustumBottom;
        for (unsigned int nbScreenshots = 0; nbScreenshots < mGridSize * mGridSize; nbScreenshots++) 
        { 
            int y = nbScreenshots / mGridSize; 
            int x = nbScreenshots - y * mGridSize; 
            
            // Shoggoth frustum extents setting
            // compute the new frustum extents
            frustumLeft        = originalFrustumLeft + frustumGridStepHorizontal * x;
            frustumRight    = frustumLeft + frustumGridStepHorizontal;
            frustumTop        = originalFrustumTop - frustumGridStepVertical * y;
            frustumBottom    = frustumTop - frustumGridStepVertical;
            
            // set the frustum extents value to the camera
            camera->setFrustumExtents(frustumLeft, frustumRight, frustumTop, frustumBottom);

            // ignore time duration between frames
            MyApplication::getSingletonPtr()->getOgreRoot()->clearEventTimes();
            mRT->update();        //render
            
            //define the current 
            Box subBox = Box(x* mWindowWidth,y * mWindowHeight,x * mWindowWidth + mWindowWidth, y * mWindowHeight + mWindowHeight);
            //copy the content from the temp buffer into the final picture PixelBox 
            //Place the tempBuffer content at the right position
            mBuffer->blitToMemory(mFinalPicturePB.getSubVolume(subBox));

        }
        
        // set frustum extents to previous settings
        camera->resetFrustumExtents();
        
        Image finalImage; //declare the final Image Object
        //insert the PixelBox data into the Image Object
        finalImage = finalImage.loadDynamicImage(static_cast<unsigned char*>(mFinalPicturePB.data), mFinalPicturePB.getWidth(),mFinalPicturePB.getHeight(),PF_B8G8R8);
        // Save the Final image to a file
        finalImage.save(fileName + "." + mFileExtension);
        
    }

    // do we have to re-enable our overlays?
    if(enableOverlayFlag)
        myApplication::getSingletonPtr()->getDefaultViewport()->setOverlaysEnabled(true);

    // reset time since last frame to pause the scene
    MyApplication::getSingletonPtr()->getOgreRoot()->clearEventTimes();
}